{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKF2S9GaYYIs"
      },
      "source": [
        "# License\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at:\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1e-w5dfNYZVg"
      },
      "source": [
        "# Instructions\n",
        "\n",
        "---\n",
        "To start an experiment:\n",
        "---\n",
        "1. Choose the agent type by setting the `AGENT_CLASS` variable in the configuration below.\n",
        "Select `Vit`  in order to choose the singlepath agent. Select `MultiVit` in order to the multipath agent.\n",
        "These agents use a ViT-Large root model.\n",
        "Set agent types suffixed with `T3` to use a ViT-Tiny root model capped to 3 layers. Set agent types suffixed with `B` to use a ViT-Base root model.\n",
        "\n",
        "1. Set `TASK_NAME` to the string-id of the task assigned to the instantiation of the selected agent.\n",
        "Refer to `TFDS_IMAGE_CLASSIFCATON_DATASETS` and `VTAB_TASKS` (below) for lists of tasks ids that have been tested with the current code. \n",
        "These lists contain task ids from the [Tensorflow Datasets Catalog](https://www.tensorflow.org/datasets/catalog/overview).\n",
        "Note that some tasks require manual download, refer to the corresponding catalog page for instructions. **WARNING**: The system state needs to be populated with at least one root model before running an agent training on any task. In order to generate the root model, set `TASK_NAME` to either `\"root_model/checkpoint\"` or `\"root_model/random_init\"` for respectively loading a pretrained root model or generating a randomly initialized one.\n",
        "\n",
        "1. Set `NUM_PATHS_SAMPLED_MAX` to the desired number of paths to sample. Additional configuration parameters can be modified in the Agents definitions code below. Configurations are set to the settings described in the publication.\n",
        "\n",
        "1. Set `SYSTEM_STATE_RELATIVE_DIR` to a relative path from where the system state will be read and written.\n",
        "\n",
        "1. By default, the system state is stored under a temporary folder within the Virtual Machine (VM) memory. This temporary folder is deleted when the VM is stopped or restarted.\n",
        "It is possible to store the system state folder on your Google Drive by activating the\n",
        "`SAVE_SYSTEM_STATE_ON_GOOGLE_DRIVE` option. In this case, you will be prompted for access approval and the system state folder will be saved in a folder named `\"munet_system_state\"` under your Google Drive root folder. Furthermore, it is also possible to store the system state into a Google Drive folder shared with multiple users by creating a link to the shared folder into your Google Drive and then setting `GDRIVE_ROOT_DIR` (below) to the path of the linked shared folder. \n",
        "\n",
        "1. To start the experiment, select \"Connect to a hosted runtime\" from the dropdown menu on the top right, and then select \"Run all\" from the \"Runtime\" menu. A free hosted CPU runtime is always available. Free access to GPU and TPU accelerators are occasionally provided by the service depending on availability.\n",
        "\n",
        "---\n",
        "During the experiment execution:\n",
        "---\n",
        "\n",
        "1. The print output is displayed after the last cell of this Colab.\n",
        "\n",
        "1. The system state folder is populated with a subfolder for each agent.\n",
        "The name of each agent folder is prefixed with the `AGENT` type string and suffixed with the `TASK_NAME`.\n",
        "Each agent directory is populated with incremental state subfolders  containing the sharded state of the architectures and parameters generated by the agent during the corresponding evolutionary cycle.\n",
        "\n",
        "1. Agents can be started asynchronously and run in parallel in varying quantities.\n",
        "It is possible to resume an interrupted agent training by restarting the execution with the same configuration.\n",
        "It is possible to continue a completed training by increasing `NUM_PATHS_SAMPLED_MAX`.\n",
        "\n",
        "1. To achieve a multi-agent execution, multiple Colabs need to be run in parallel, each set to the same configuration but different `TASK_NAME`.\n",
        "\n",
        "1. To achieve heterogeneous hardware execution, parallel Colab Notebooks can be connected to a runtime of different types.\n",
        "It is possible to switch between CPU, GPU and TPU by selecting `Change runtime type` in the `Resources` tab in this Colab Notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M93tll7z29rX"
      },
      "outputs": [],
      "source": [
        "# @title Agent parameters\n",
        "AGENT_CLASS = \"VitT3\" # @param [\"VitT3\", \"VitB\", \"Vit\", \"MultiVitT3\", \"MultiVitB\", \"MultiVit\"] { type: \"string\", isTemplate: true }\n",
        "# Set TASK_NAME to \"root_model/checkpoint\" or \"root_model/random_init\" to initalize the population.\n",
        "TASK_NAME = \"root_model/checkpoint\"  # @param { type: \"string\", isTemplate: true }\n",
        "NUM_PATHS_SAMPLED_MAX = 16 # @param { type: \"integer\", isTemplate: true }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xcaI2KII_MH5"
      },
      "outputs": [],
      "source": [
        "# Saves system state on Google drive instead of saving it in a temporary VM folder.\n",
        "SYSTEM_STATE_RELATIVE_DIR = \"munet_system_state/\"  # @param { type: \"string\", isTemplate: true }\n",
        "SAVE_SYSTEM_STATE_ON_GOOGLE_DRIVE = False  # @param { type: \"boolean\", isTemplate: true }\n",
        "if SAVE_SYSTEM_STATE_ON_GOOGLE_DRIVE:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/gdrive')\n",
        "  GDRIVE_ROOT_DIR = \"/content/gdrive/My Drive/\"\n",
        "  SYSTEM_STATE_DIR = GDRIVE_ROOT_DIR + SYSTEM_STATE_RELATIVE_DIR\n",
        "  print(\"Saving system state in Google Drive.\")\n",
        "else:\n",
        "  SYSTEM_STATE_DIR = \"/tmp/\" + SYSTEM_STATE_RELATIVE_DIR\n",
        "  print(\"WARNING: Saving system state in VM, state will be lost after reboot!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dsiDs_mgBZOx"
      },
      "outputs": [],
      "source": [
        "# Test immutability of published paths at beginning of each cycle.\n",
        "# Tollerance may be increased if the system is run with any context difference: e.g. harware, input preprocessing, libraries or datsets version.\n",
        "TEST_IMMUTABILITY = False\n",
        "IMMUTABILITY_RELATIVE_TOLLERANCE = 0.001  # 0.1%"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KZ0njfC-XCBA"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SdaJ0somYxIK"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade -q pip jax jaxlib\n",
        "!pip install --upgrade -q git+https://github.com/google/flax.git\n",
        "!pip install -q ml_collections\n",
        "!pip install -q tensorflow_addons\n",
        "![ -d task_adaptation ] || git clone --depth=1 https://github.com/google-research/task_adaptation\n",
        "![ -d vision_transformer ] || git clone --depth=1 https://github.com/google-research/vision_transformer\n",
        "\n",
        "import sys\n",
        "if './task_adaptation' not in sys.path:\n",
        "  sys.path.append('./task_adaptation')\n",
        "if './vision_transformer' not in sys.path:\n",
        "  sys.path.append('./vision_transformer')\n",
        "\n",
        "import jax.tools.colab_tpu\n",
        "try:\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "except:\n",
        "  pass  # Not a Tpu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LzhEweKzMg6k"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import flax\n",
        "import flax.linen as nn\n",
        "import gc\n",
        "import inspect\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import json\n",
        "import math\n",
        "import numpy as np\n",
        "import optax\n",
        "import os\n",
        "import pandas as pd\n",
        "import random\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import time\n",
        "from collections import defaultdict\n",
        "from functools import partial\n",
        "from flax.training import checkpoints as flax_checkpoints\n",
        "from ml_collections import ConfigDict, FrozenConfigDict\n",
        "from tensorflow.io import gfile\n",
        "from threading import Thread, Lock\n",
        "from typing import Any, Type\n",
        "tf.compat.v1.enable_eager_execution()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6QHBzcuUYeh5"
      },
      "outputs": [],
      "source": [
        "# ViT imports.\n",
        "from vision_transformer.vit_jax import checkpoint as vit_checkpoint\n",
        "from vision_transformer.vit_jax.configs import models as vit_models_config\n",
        "from vision_transformer.vit_jax import models_vit as vit_models\n",
        "# VTAB imports.\n",
        "import task_adaptation.registry as task_adapt_registry\n",
        "import task_adaptation.data.caltech\n",
        "import task_adaptation.data.cifar\n",
        "import task_adaptation.data.dtd\n",
        "import task_adaptation.data.oxford_flowers102\n",
        "import task_adaptation.data.oxford_iiit_pet\n",
        "import task_adaptation.data.sun397\n",
        "import task_adaptation.data.svhn\n",
        "import task_adaptation.data.patch_camelyon\n",
        "import task_adaptation.data.eurosat\n",
        "import task_adaptation.data.resisc45\n",
        "import task_adaptation.data.diabetic_retinopathy\n",
        "import task_adaptation.data.clevr\n",
        "import task_adaptation.data.dmlab\n",
        "import task_adaptation.data.dsprites\n",
        "import task_adaptation.data.kitti\n",
        "import task_adaptation.data.smallnorb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zhl1L8ldXFJz"
      },
      "source": [
        "# Utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SHBWj0JmpWDX"
      },
      "outputs": [],
      "source": [
        "# Ref. Tfds catalog: https://www.tensorflow.org/datasets/catalog/overview\n",
        "TFDS_IMAGE_CLASSIFCATON_TASKS = {\n",
        "    \"beans\": {\"version\":\"0.1.0\"},\n",
        "    \"binary_alpha_digits\": {\"version\":\"1.0.0\"},\n",
        "    \"caltech_birds2010\": {\"version\":\"0.1.1\"},\n",
        "    \"caltech_birds2011\": {\"version\":\"0.1.1\"},\n",
        "    \"cars196\": {\"version\":\"2.1.0\"},\n",
        "    \"cassava\": {\"version\":\"0.1.0\"},\n",
        "    \"cats_vs_dogs\": {\"version\":\"4.0.0\"},\n",
        "    \"cifar10\": {\"version\":\"3.0.2\"},\n",
        "    \"cifar100\": {\"version\":\"3.0.2\"},\n",
        "    \"citrus_leaves\": {\"version\":\"0.1.2\"},\n",
        "    \"cmaterdb/bangla\": {\"version\":\"1.0.0\"},\n",
        "    \"cmaterdb/devanagari\": {\"version\":\"1.0.0\"},\n",
        "    \"cmaterdb/telugu\": {\"version\":\"1.0.0\"},\n",
        "    \"colorectal_histology\": {\"version\":\"2.0.0\"},\n",
        "    \"controlled_noisy_web_labels/mini_imagenet_red\": {\"version\":\"1.0.0\"},\n",
        "    \"controlled_noisy_web_labels/mini_imagenet_blue\": {\"version\":\"1.0.0\"},\n",
        "    \"curated_breast_imaging_ddsm/patches\": {\"version\":\"3.0.0\"},\n",
        "    \"cycle_gan/apple2orange\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/summer2winter_yosemite\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/horse2zebra\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/monet2photo\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/cezanne2photo\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/ukiyoe2photo\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/vangogh2photo\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/maps\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/cityscapes\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/facades\": {\"version\":\"2.0.0\"},\n",
        "    \"cycle_gan/iphone2dslr_flower\": {\"version\":\"2.0.0\"},\n",
        "    \"deep_weeds\": {\"version\":\"3.0.0\"},\n",
        "    \"domainnet/real\": {\"version\":\"1.0.0\"},\n",
        "    \"domainnet/painting\": {\"version\":\"1.0.0\"},\n",
        "    \"domainnet/clipart\": {\"version\":\"1.0.0\"},\n",
        "    \"domainnet/quickdraw\": {\"version\":\"1.0.0\"},\n",
        "    \"domainnet/infograph\": {\"version\":\"1.0.0\"},\n",
        "    \"domainnet/sketch\": {\"version\":\"1.0.0\"},\n",
        "    \"emnist/digits\": {\"version\":\"3.0.0\"},\n",
        "    \"emnist/letters\": {\"version\":\"3.0.0\"},\n",
        "    \"fashion_mnist\": {},  # Not yet included in the µNet system checkpoint.\n",
        "    \"food101\": {\"version\":\"2.0.0\"},\n",
        "    \"horses_or_humans\": {\"version\":\"3.0.0\"},\n",
        "    \"i_naturalist2017\": {\"version\":\"0.1.0\"},\n",
        "    \"i_naturalist2018\": {\"version\":\"1.0.0\"},\n",
        "    \"i_naturalist2021\": {},  # Not yet included in the µNet system checkpoint.\n",
        "    \"imagenet2012\": {\"version\":\"5.1.0\"},\n",
        "    \"imagenet_a\": {\"version\":\"0.1.0\"},\n",
        "    \"imagenet_lt\": {\"version\":\"1.0.0\"},\n",
        "    \"imagenet_r\": {\"version\":\"0.2.0\"},\n",
        "    \"imagenet_sketch\": {\"version\":\"1.0.0\"},\n",
        "    \"imagenette\": {\"version\":\"1.0.0\"},\n",
        "    \"imagewang\": {\"version\":\"2.0.0\"},\n",
        "    \"kmnist\": {\"version\":\"3.0.1\"},\n",
        "    \"malaria\": {\"version\":\"1.0.0\"},\n",
        "    \"mnist\": {\"version\":\"3.0.1\"},\n",
        "    \"omniglot\": {\"version\":\"3.0.0\"},\n",
        "    \"pet_finder\": {\"version\":\"1.0.0\"},\n",
        "    \"places365_small\": {\"version\":\"2.1.0\"},\n",
        "    \"plant_village\": {\"version\":\"1.0.2\"},\n",
        "    \"plantae_k\": {\"version\":\"0.1.0\"},\n",
        "    \"quickdraw_bitmap\": {\"version\":\"3.0.0\"},\n",
        "    \"rock_paper_scissors\": {\"version\":\"3.0.0\"},\n",
        "    \"siscore/rotation\": {\"version\":\"1.0.0\"},\n",
        "    \"siscore/size\": {\"version\":\"1.0.0\"},\n",
        "    \"siscore/location\": {\"version\":\"1.0.0\"},\n",
        "    \"stanford_dogs\": {\"version\":\"0.2.0\"},\n",
        "    \"stanford_online_products\": {\"version\":\"1.0.0\"},\n",
        "    \"stl10\": {\"version\":\"1.0.0\"},\n",
        "    \"tf_flowers\": {\"version\":\"3.0.1\"},\n",
        "    \"uc_merced\": {\"version\":\"2.0.0\"},\n",
        "    \"visual_domain_decathlon/aircraft\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/cifar100\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/daimlerpedcls\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/dtd\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/gtsrb\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/imagenet12\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/omniglot\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/svhn\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/ucf101\": {\"version\":\"1.2.0\"},\n",
        "    \"visual_domain_decathlon/vgg-flowers\": {\"version\":\"1.2.0\"},\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XI7xcfWLXx-G"
      },
      "outputs": [],
      "source": [
        "# Append suffix \"/1k\" to get the 1k version of each task.\n",
        "VTAB_TASKS = [\n",
        "    \"caltech101\",\n",
        "    # cifar100/10 were already added with slightly different val split but same test set. So here is added only the 1k versions.\n",
        "    \"cifar100/1k\",\n",
        "    \"cifar10/1k\",\n",
        "    \"dtd\",\n",
        "    \"oxford_flowers102\",\n",
        "    \"oxford_iiit_pet\",\n",
        "    \"sun397\",\n",
        "    \"svhn_cropped\",\n",
        "    \"patch_camelyon\",\n",
        "    \"eurosat\",\n",
        "    \"resisc45\",\n",
        "    \"diabetic_retinopathy_detection/btgraham-300\",\n",
        "    \"clevr/count_cylinders\",  # Not in results table.\n",
        "    \"clevr/count_all\",  # Clevr-Count\n",
        "    \"clevr/closest_object_distance\",  # Clevr-Dist\n",
        "    \"dmlab\",\n",
        "    \"dsprites/label_x_position\",  # dSpr-Loc\n",
        "    \"dsprites/label_orientation\",  # dSpr-Ori\n",
        "    \"kitti/closest_object_distance\",  # Not in results table.\n",
        "    \"kitti/count_vehicles\",  # Not in results table.\n",
        "    \"kitti/closest_vehicle_distance\",  # Kitti-dist\n",
        "    \"smallnorb/label_category\",  # Not in results table.\n",
        "    \"smallnorb/label_lighting\",  # Not in results table.\n",
        "    \"smallnorb/label_azimuth\",  # Azim\n",
        "    \"smallnorb/label_elevation\",  # Elev\n",
        "    ]\n",
        "for tn in VTAB_TASKS:\n",
        "  assert tn not in TFDS_IMAGE_CLASSIFCATON_TASKS.keys(), tn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OeLYCqhfXZyk"
      },
      "outputs": [],
      "source": [
        "def compute_flops_hlo(flax_module, *a, **kw):\n",
        "  m = jax.xla_computation(flax_module)(*a, **kw).as_hlo_module()\n",
        "  # Compute flops on cpu for cross platform concistency.\n",
        "  client = jax.lib.xla_bridge.get_backend(\"cpu\")\n",
        "  analysis = jax.lib.xla_client._xla.hlo_module_cost_analysis(client, m)\n",
        "  return analysis[\"flops\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AgTEwq24TYn8"
      },
      "source": [
        "# Models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ueAYFssCUG78"
      },
      "outputs": [],
      "source": [
        "def get_sample_images(image_size, batch_size):\n",
        "  return np.zeros((batch_size, image_size, image_size, 3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tkv4dVcQV4YB"
      },
      "outputs": [],
      "source": [
        "def get_num_params(params):\n",
        "  return sum(jax.tree_util.tree_flatten(\n",
        "      jax.tree_util.tree_map(lambda p: np.prod(p.shape), params))[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V6vIOQBEVWch"
      },
      "outputs": [],
      "source": [
        "def get_optimizer(\n",
        "    opt_lr: float,\n",
        "    opt_lr_schedule: str,\n",
        "    opt_lr_warmup_ratio: float,\n",
        "    opt_momentum: float,\n",
        "    opt_nesterov: bool,\n",
        "    num_train_batches_between_validations: int,\n",
        "    num_validations_per_path_training: int,\n",
        "    ):\n",
        "  min_lr = opt_lr / 1000.0\n",
        "  if opt_lr_schedule == \"constant\":\n",
        "    # Divide by 2 so that average lr is the same as other types.\n",
        "    learning_rate = 0.5 * opt_lr\n",
        "  elif opt_lr_schedule == \"linear\":\n",
        "    train_steps = int(num_train_batches_between_validations * num_validations_per_path_training)\n",
        "    warmup_steps = int(opt_lr_warmup_ratio * train_steps)\n",
        "    schedules = [\n",
        "        optax.linear_schedule(\n",
        "            init_value=min_lr,\n",
        "            end_value=opt_lr,\n",
        "            transition_steps=warmup_steps),\n",
        "        optax.linear_schedule(\n",
        "            init_value=opt_lr,\n",
        "            end_value=min_lr,\n",
        "            transition_steps=train_steps-warmup_steps)]\n",
        "    learning_rate = optax.join_schedules(schedules, [warmup_steps])\n",
        "  elif opt_lr_schedule == \"cosine\":\n",
        "    train_steps = int(num_train_batches_between_validations\n",
        "                      * num_validations_per_path_training)\n",
        "    learning_rate = optax.warmup_cosine_decay_schedule(\n",
        "        init_value=min_lr,\n",
        "        peak_value=opt_lr,\n",
        "        warmup_steps=int(opt_lr_warmup_ratio * train_steps),\n",
        "        decay_steps=train_steps)\n",
        "  elif opt_lr_schedule == \"restarts\":\n",
        "    train_steps = num_train_batches_between_validations\n",
        "    repeats = num_validations_per_path_training\n",
        "    kwargs = dict(\n",
        "        init_value=min_lr,\n",
        "        peak_value=opt_lr,\n",
        "        warmup_steps=int(opt_lr_warmup_ratio * train_steps),\n",
        "        decay_steps=train_steps,\n",
        "    )\n",
        "    kwargs = [kwargs] * repeats\n",
        "    learning_rate = optax.sgdr_schedule(kwargs)\n",
        "  else:\n",
        "    assert False, f\"Invalid lr schedule: {opt_lr_schedule}\"\n",
        "\n",
        "  return optax.chain(\n",
        "      optax.clip_by_global_norm(1.0),\n",
        "      optax.sgd(\n",
        "          learning_rate=learning_rate,\n",
        "          momentum=opt_momentum,\n",
        "          nesterov=opt_nesterov,\n",
        "          accumulator_dtype=jnp.bfloat16))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2GpeOYxFi1Sm"
      },
      "outputs": [],
      "source": [
        "def merge_params(a, b):\n",
        "  params = a.copy(b)\n",
        "  assert len(params) == len(a) + len(b)\n",
        "  return params"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ExYTVeamDcPy"
      },
      "source": [
        "## Vit Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "43JOxcGVMJ4s"
      },
      "outputs": [],
      "source": [
        "class VitModelFactory():\n",
        "  @staticmethod\n",
        "  def get_model(hparams, config):\n",
        "    return get_vit_model(hparams, config)\n",
        "\n",
        "  @staticmethod\n",
        "  def get_init_comps(hparams, config):\n",
        "    return get_vit_init_comps(hparams, config)\n",
        "\n",
        "  @staticmethod\n",
        "  def get_comps2model_fn():\n",
        "    return vit_comps2model\n",
        "\n",
        "  @staticmethod\n",
        "  def get_sample_input(hparams):\n",
        "    return get_sample_images(image_size=hparams[\"ds_image_size\"], batch_size=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vvZ_4-kJ9Pt3"
      },
      "outputs": [],
      "source": [
        "def get_vit_filename(query):\n",
        "  df = vit_checkpoint.get_augreg_df()\n",
        "  res = df.query(query).filename.unique()\n",
        "  assert len(res) == 1\n",
        "  return res[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0lvqd47g9ZsW"
      },
      "outputs": [],
      "source": [
        "VIT_CONFIG_CACHE = {}\n",
        "\n",
        "def get_vit_config(query):\n",
        "  global VIT_CONFIG_CACHE\n",
        "  if query not in VIT_CONFIG_CACHE:\n",
        "    filename = get_vit_filename(query)\n",
        "    config = vit_models_config.AUGREG_CONFIGS[filename.split(\"-\")[0]].copy_and_resolve_references()\n",
        "    config.unlock()\n",
        "    # Disable dropout.\n",
        "    config.transformer.dropout_rate = 0.0\n",
        "    config.transformer.attention_dropout_rate = 0.0\n",
        "    config.lock()\n",
        "    VIT_CONFIG_CACHE[query] = config\n",
        "  return VIT_CONFIG_CACHE[query].copy_and_resolve_references()\n",
        "\n",
        "def get_set_vit_config(hparams, config):\n",
        "  path_config = get_vit_config(config.vit_checkpoint_query)\n",
        "  path_config.transformer.num_layers = int(hparams[\"num_layers\"])\n",
        "  path_config.unlock()\n",
        "  path_config.num_classes = int(hparams[\"num_classes\"])\n",
        "  if \"classifier\" in hparams:\n",
        "    path_config.classifier = hparams[\"classifier\"]\n",
        "  path_config.lock()\n",
        "  path_config = FrozenConfigDict(path_config)\n",
        "  return path_config\n",
        "\n",
        "def get_max_num_layers(query):\n",
        "  config = get_vit_config(query)\n",
        "  return config.transformer.num_layers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JOV3CPHMUQbG"
      },
      "outputs": [],
      "source": [
        "# Get params from ViT checkpoints.\n",
        "def get_vit_checkpoint_comps(image_size, query):\n",
        "  filename = get_vit_filename(query)\n",
        "  config = get_vit_config(query)\n",
        "  model = vit_models.VisionTransformer(**config, num_classes=1)  # num_classes unused.\n",
        "  init_params = copy.deepcopy(jax.device_get(\n",
        "      model.init(jax.random.PRNGKey(random.randrange(int(1e10))),\n",
        "                 VitModelFactory.get_sample_input({\"ds_image_size\": image_size}),\n",
        "                 train=False  # Disables dropout, no effect on params.\n",
        "                 )[\"params\"]))\n",
        "  params = vit_checkpoint.load_pretrained(\n",
        "    pretrained_path=f\"gs://vit_models/augreg/{filename}.npz\",\n",
        "    init_params=init_params,\n",
        "    model_config=config)\n",
        "  return vit_model2comps(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FjEWZgFIUt25"
      },
      "outputs": [],
      "source": [
        "# Get ViT model and init_params.\n",
        "def get_vit_model(hparams, config):\n",
        "  vit_config = get_set_vit_config(hparams, config)\n",
        "  return vit_models.VisionTransformer(**vit_config)\n",
        "\n",
        "def get_vit_init_comps(hparams, config):\n",
        "  model = get_vit_model(hparams, config)\n",
        "  init_params = copy.deepcopy(jax.device_get(model.init(\n",
        "      jax.random.PRNGKey(random.randrange(int(1e10))),\n",
        "      VitModelFactory.get_sample_input(hparams),\n",
        "      train=False  # Disables dropout, no effect on params.\n",
        "      )[\"params\"]))\n",
        "  return vit_model2comps(init_params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "txU4go5GUf-e"
      },
      "outputs": [],
      "source": [
        "# ViT parameters mapping to components.\n",
        "TRANSFORMER_KEYS = set(\n",
        "    [\"encoder_norm\", \"posembed_input\" ] + \\\n",
        "    [f\"encoderblock_{k}\" for k in range(30)])\n",
        "\n",
        "def vit_model2comps(params):\n",
        "  new_params = {}\n",
        "  for k in params.keys():\n",
        "    if k == \"Transformer\":\n",
        "      t_params = params[k]\n",
        "      for t_k in t_params.keys():\n",
        "        new_params[t_k] = t_params[t_k]\n",
        "    else:\n",
        "      new_params[k] = params[k]\n",
        "  return flax.core.freeze(new_params)\n",
        "\n",
        "def vit_comps2model(params):\n",
        "  new_params = params.unfreeze()\n",
        "  new_params[\"Transformer\"] = {}\n",
        "  for k in list(new_params.keys()):\n",
        "    if k in TRANSFORMER_KEYS:\n",
        "      new_params[\"Transformer\"][k] = new_params.pop(k)\n",
        "  assert len(new_params[\"Transformer\"]) != 0\n",
        "  return flax.core.freeze(new_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ShsZ14HoiuFD"
      },
      "source": [
        "## MultiVit Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z0-D1trkMOor"
      },
      "outputs": [],
      "source": [
        "class MultiVitModelFactory():\n",
        "  @staticmethod\n",
        "  def get_model(hparams, config):\n",
        "    return get_multivit_model(hparams, config)\n",
        "\n",
        "  @staticmethod\n",
        "  def get_init_comps(hparams, config):\n",
        "    return get_multivit_init_comps(hparams, config)\n",
        "\n",
        "  @staticmethod\n",
        "  def get_comps2model_fn():\n",
        "    return multivit_comps2model\n",
        "\n",
        "  @staticmethod\n",
        "  def get_sample_input(hparams):\n",
        "    return {str(k): get_sample_images(image_size=k, batch_size=1) for k in hparams[\"ds_image_size\"]}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "61WJRSygqedz"
      },
      "outputs": [],
      "source": [
        "def get_multivit_init_comps(hparams, config):\n",
        "  model = get_multivit_model(hparams, config)\n",
        "  init_params = copy.deepcopy(jax.device_get(\n",
        "      model.init(\n",
        "          jax.random.PRNGKey(random.randrange(int(1e10))),\n",
        "          MultiVitModelFactory.get_sample_input(hparams),\n",
        "          train=False  # Disables dropout, no effect on params.\n",
        "          )[\"params\"]))\n",
        "  return multivit_model2comps(init_params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YeB-D_HVdIDm"
      },
      "outputs": [],
      "source": [
        "def multivit_comps2model(params):\n",
        "  params = params.unfreeze()\n",
        "  for k in params:\n",
        "    if k.startswith(\"path_\"):\n",
        "      params[k] = vit_comps2model(flax.core.freeze(params[k]))\n",
        "  return flax.core.freeze(params)\n",
        "\n",
        "def multivit_model2comps(params):\n",
        "  # Mapping of paths component skipped since those are never used from rand init.\n",
        "  return params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "waN2hnq7kbOC"
      },
      "outputs": [],
      "source": [
        "class MultipathRouter(nn.Module):\n",
        "  init_main_path_weight: float\n",
        "  num_paths: int\n",
        "  lr_mult: float\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    assert self.num_paths \u003e 0\n",
        "    assert self.lr_mult \u003e= 0 and self.lr_mult \u003c= 1\n",
        "    init_bias = np.log((1/(1/self.init_main_path_weight -1))*(self.num_paths-1))\n",
        "    x = nn.LayerNorm()(x)\n",
        "    x = nn.Dense(self.num_paths,\n",
        "                 kernel_init=nn.initializers.zeros,\n",
        "                 bias_init=nn.initializers.constant(np.asarray([init_bias]+[0]*(self.num_paths-1)))\n",
        "                 )(x)\n",
        "    x = nn.softmax(x)\n",
        "    x = self.lr_mult * x + (1-self.lr_mult) * jax.lax.stop_gradient(x)\n",
        "    return x\n",
        "\n",
        "class Connector(nn.Module):\n",
        "  out_dim: int\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    x = nn.Dense(self.out_dim,\n",
        "                 kernel_init=nn.initializers.zeros,\n",
        "                 bias_init=nn.initializers.zeros,\n",
        "                 )(x)\n",
        "    return x\n",
        "\n",
        "class MultiVitModel(nn.Module):\n",
        "  config: Any\n",
        "  path_module: Type[nn.Module] = vit_models.VisionTransformer\n",
        "  main_path_name: str = \"path_0\"\n",
        "  @nn.compact\n",
        "  def __call__(self, inputs, *, train):\n",
        "    logits_0 = self.path_module(\n",
        "        name=self.main_path_name,\n",
        "        **self.config.paths_configs[self.main_path_name])(\n",
        "            inputs[self.config.paths_image_size[self.main_path_name]], train=train)\n",
        "    out_dim = logits_0.shape[-1]\n",
        "    weights = MultipathRouter(\n",
        "        name=\"multipath_router\",\n",
        "        num_paths=len(self.config.paths_configs),\n",
        "        **self.config.router)(\n",
        "            logits_0)\n",
        "    all_logits = [logits_0]\n",
        "    for path_name in self.config.paths_configs:\n",
        "      if path_name == self.main_path_name:\n",
        "        continue\n",
        "      representation = self.path_module(\n",
        "          name=path_name,\n",
        "          **self.config.paths_configs[path_name])(\n",
        "              inputs[self.config.paths_image_size[path_name]], train=train)\n",
        "      path_logits = Connector(\n",
        "          name=f\"head_adapter_{path_name}\", out_dim=out_dim)(representation)\n",
        "      all_logits.append(path_logits)\n",
        "    stacked = jnp.stack(all_logits, axis=-1)\n",
        "    logits_comb = jnp.einsum(\"BLp,Bp-\u003eBL\", stacked, weights)\n",
        "    logits_comb = jnp.einsum(\"BLp,Bp-\u003eBL\", jax.lax.stop_gradient(stacked), weights)\n",
        "    logits_sum = jnp.einsum(\"BLp,Bp-\u003eBL\", stacked, jnp.ones_like(weights))\n",
        "    logits_out = logits_comb - jax.lax.stop_gradient(logits_sum) + logits_sum\n",
        "    return logits_out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CC4pF48eqDLL"
      },
      "outputs": [],
      "source": [
        "def get_multivit_model(hparams, config):\n",
        "  model_config = ConfigDict()\n",
        "  model_config.paths_configs = {\n",
        "      k: get_set_vit_config(hparams[\"paths\"][k][\"hparams\"], config) for k in hparams[\"paths\"]}\n",
        "  model_config.paths_image_size = {\n",
        "      k: str(hparams[\"paths\"][k][\"hparams\"][\"ds_image_size\"]) for k in hparams[\"paths\"]}\n",
        "  model_config.router = {\n",
        "      \"init_main_path_weight\": float(hparams[\"router_init_main_path_weight\"]),\n",
        "      \"lr_mult\": float(hparams[\"router_lr_mult\"]),\n",
        "  }\n",
        "  return MultiVitModel(config=FrozenConfigDict(model_config))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lyX6WjaUDWOi"
      },
      "source": [
        "# Agents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FNAi2xK-fARf"
      },
      "outputs": [],
      "source": [
        "def format_agent_id(class_name, task_name):\n",
        "  agent_id = f\"{class_name}/{task_name}\"\n",
        "  assert \"~\" not in agent_id, f\"Invalid agent id: {agent_id}\"\n",
        "  return agent_id.replace(\"/\", \"~\")\n",
        "\n",
        "def get_agent_class(agent_id):\n",
        "  return globals()[agent_id.split(\"~\")[0]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TgGHyptb_BwQ"
      },
      "outputs": [],
      "source": [
        "def incremental_mutation(value, values_list):\n",
        "  assert value in values_list, f\"{value} not in {values_list}\"\n",
        "  idx = values_list.index(value)\n",
        "  idx += 1 if np.random.uniform() \u003c 0.5 else -1\n",
        "  idx = max(0, min(len(values_list)-1, idx))\n",
        "  return values_list[idx]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_dN3bJiaODsv"
      },
      "outputs": [],
      "source": [
        "def decay_selection(paths):\n",
        "  for path in sorted(paths, key=lambda p: p.score(), reverse=True):\n",
        "    offsprings = path.metrics.get(\"offsprings\", 0)\n",
        "    assert not math.isnan(offsprings)\n",
        "    select_prob = 0.5 ** offsprings\n",
        "    print(f\"  Candidate path {path.id},\",\n",
        "          f\"selection probability: 0.5^{offsprings}=={select_prob}\")\n",
        "    if np.random.uniform() \u003c select_prob:\n",
        "      print(f\"  Selected path {path.id}\")\n",
        "      return path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jCNLTP1zGu_u"
      },
      "outputs": [],
      "source": [
        "class Agent():\n",
        "  def __init__(self, system_state_dir, task_name, num_paths_sampled_max):\n",
        "    self.config = self.get_config()\n",
        "    self.config.system_state_dir = system_state_dir\n",
        "    self.config.task_name = task_name\n",
        "    self.config.agent_id = format_agent_id(self.class_name, task_name)\n",
        "    self.config.agent_dir = os.path.join(system_state_dir, self.id)\n",
        "    self.config.num_paths_sampled_max = num_paths_sampled_max\n",
        "    self.config = FrozenConfigDict(self.config)\n",
        "    config_validate(self.config)\n",
        "    self.task = None if task_name.startswith(\"root_model/\") else Task(name=task_name, config=self.config)\n",
        "\n",
        "  @property\n",
        "  def class_name(self):\n",
        "    return self.__class__.__name__\n",
        "\n",
        "  @property\n",
        "  def id(self):\n",
        "    return self.config.agent_id\n",
        "\n",
        "  def agents_to_load(self):\n",
        "    # Defaults to load only agents of the same class. Extend this list to allow\n",
        "    # to access the strucutres and parameters produced by different agent types.\n",
        "    return [self.class_name+\"~*\"]\n",
        "\n",
        "  def prune_population(self):\n",
        "    print(\"Prune population\")\n",
        "    candidates = copy.copy(self.pop.paths[self.config.agent_id])\n",
        "    survivors = []\n",
        "    while candidates and len(survivors) \u003c self.config.max_population_size:\n",
        "      print(\"Sampling survivor:\", len(survivors)+1)\n",
        "      selected = decay_selection(candidates) if survivors else None\n",
        "      if selected is None:\n",
        "        selected = get_best_path(candidates)\n",
        "        print(\"  Selected best candidate:\", selected.id)\n",
        "      assert selected is not None\n",
        "      survivors.append(selected)\n",
        "      candidates.remove(selected)\n",
        "    assert len(survivors) == len(set(survivors))  # Check no duplicates.\n",
        "    self.pop.paths[self.config.agent_id] = survivors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5j4LPnTA8EzH"
      },
      "outputs": [],
      "source": [
        "def run_evolution(agent):\n",
        "  config = agent.config\n",
        "  devices = jax.local_devices()\n",
        "  print(\"DEVICE COUNT:\", len(devices))\n",
        "  while True:\n",
        "    print(\"\\n\\n====\")\n",
        "    agent.load_state()\n",
        "    print(f\"\\nSAMPLES: [{agent.num_paths_sampled}/{config.num_paths_sampled_max}]\")\n",
        "    if agent.num_paths_sampled \u003e= config.num_paths_sampled_max:\n",
        "      break\n",
        "    task = agent.task\n",
        "    best_path = agent.pop.get_best_path()\n",
        "    if TEST_IMMUTABILITY and best_path:\n",
        "      run_test_eval(best_path, task, test_immutability=True)\n",
        "    gen_hparams = agent.sample_gen_hps()\n",
        "    gen_hparams[\"num_classes\"] = task.num_classes\n",
        "    paths = []\n",
        "    for i in range(len(devices)):\n",
        "      print(f\"Sampling path: {Path.counter}\")\n",
        "      paths.append(agent.sample_path(gen_hparams))\n",
        "      agent.num_paths_sampled += 1\n",
        "      gc.collect()\n",
        "    gen_hparams = agent.finalize_hps(gen_hparams, paths)\n",
        "    ds_train = task.get_ds(\"train\", gen_hparams)\n",
        "    ds_validation = task.get_ds(\"validation\", gen_hparams)\n",
        "    train_loop(agent, paths, ds_train, ds_validation, devices)\n",
        "    for path in paths:\n",
        "      if path.metrics[\"improved\"]:\n",
        "        assert path not in agent.pop.paths[config.agent_id]\n",
        "        agent.pop.paths[config.agent_id].append(path)\n",
        "    # Track best path.\n",
        "    curr_best_path = agent.pop.get_best_path()\n",
        "    if curr_best_path != best_path:\n",
        "      if best_path:\n",
        "        assert curr_best_path.score() \u003e= best_path.score()\n",
        "      best_path = curr_best_path\n",
        "      best_path.metrics[\"new_best\"] = True\n",
        "      agent.print_best_path_summary()\n",
        "    df_leaderboard(pop_to_df(agent.pop))\n",
        "    agent.prune_population()\n",
        "    assert best_path in agent.pop.paths[config.agent_id], best_path\n",
        "    best_path.metrics[\"num_cycles\"] = best_path.metrics.get(\"num_cycles\", 0) + 1\n",
        "    run_test_eval(best_path, task)\n",
        "    save_state(agent)\n",
        "    if agent.num_paths_sampled \u003e= config.num_paths_sampled_max:\n",
        "      break\n",
        "\n",
        "def run_root_model(agent):\n",
        "  agent.load_state()\n",
        "  save_state(agent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9I56DcbI79M2"
      },
      "source": [
        "## Vit Agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "njMI52JMnqq1"
      },
      "outputs": [],
      "source": [
        "def get_common_config_vit():\n",
        "  config = ConfigDict()\n",
        "  config.num_train_examples_between_validations_max = 300_000\n",
        "  config.num_validations_per_path_training = 4\n",
        "  config.num_validation_examples_max = 10_000\n",
        "  config.max_population_size = 4\n",
        "  # Force finetune last layer norm that technically is part of the head.\n",
        "  config.force_mutations = [\"clone:encoder_norm\"]\n",
        "  config.scorer_kwargs = dict(\n",
        "      scale_factor=0.99,\n",
        "      base_accounted_params=2_200_000_000,\n",
        "      base_flops=3_800_000_000_000,\n",
        "      )\n",
        "  config.hparams_defaults = {\n",
        "      \"_mu_\": 0.2,\n",
        "      \"opt_lr\": 0.02,\n",
        "      \"opt_lr_schedule\": \"cosine\",\n",
        "      \"opt_lr_warmup_ratio\": 0.02,\n",
        "      \"opt_momentum\": 0.8,\n",
        "      \"opt_nesterov\": True,\n",
        "      \"ds_area_range_min\": 1.0,\n",
        "      \"ds_aspect_ratio_range_min\": 1.0,\n",
        "      \"ds_flip_left_right\": False,\n",
        "      \"ds_brightness_delta\": 0.0,\n",
        "      \"ds_contrast_delta\": 0.0,\n",
        "      \"ds_saturation_delta\": 0.0,\n",
        "      \"ds_hue_delta\": 0.0,\n",
        "      \"ds_quality_delta\": 0.0,\n",
        "  }\n",
        "  config.hparams_mutation_ranges = {\n",
        "      \"_mu_\": [0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.22, 0.24, 0.26, 0.28, 0.30],\n",
        "      \"opt_lr\": [0.0001, 0.0002, 0.0005, 0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0],\n",
        "      \"opt_lr_warmup_ratio\": [0.0, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3],\n",
        "      \"opt_momentum\": [0.5, 0.6, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 0.98, 0.99],\n",
        "      \"opt_nesterov\": [True, False],\n",
        "      \"ds_area_range_min\": [0.05, 0.5, 0.95, 1.0],\n",
        "      \"ds_aspect_ratio_range_min\": [0.5, 0.75, 1.0],\n",
        "      \"ds_flip_left_right\": [True, False],\n",
        "      \"ds_brightness_delta\": [0.0, 0.01, 0.02, 0.05, 0.1, 0.2],\n",
        "      \"ds_contrast_delta\": [0.0, 0.01, 0.02, 0.05, 0.1, 0.2],\n",
        "      \"ds_saturation_delta\": [0.0, 0.01, 0.02, 0.05, 0.1, 0.2],\n",
        "      \"ds_hue_delta\": [0.0, 0.01, 0.02, 0.05, 0.1, 0.2],\n",
        "      \"ds_quality_delta\": [ 0.0, 0.01, 0.02, 0.05, 0.1, 0.2],\n",
        "  }\n",
        "  return config\n",
        "\n",
        "def get_config_vit_large():\n",
        "  config = get_common_config_vit()\n",
        "  config.batch_size = 16\n",
        "  # The query is used to get the model configs even if the checkpoint is not loaded.\n",
        "  config.vit_checkpoint_query = 'name==\"L/16\" and ds==\"i21k\" and aug==\"medium2\" and wd==0.03 and sd==0.1'\n",
        "  max_num_layers = get_max_num_layers(config.vit_checkpoint_query)\n",
        "  config.hparams_defaults[\"num_layers\"] = max_num_layers\n",
        "  config.hparams_mutation_ranges[\"num_layers\"] = list(\n",
        "      range(config.hparams_defaults[\"num_layers\"]+1\n",
        "            +1  # Allow to exceed root-model's layers by 1.\n",
        "            ))\n",
        "  config.hparams_defaults[\"ds_image_size\"] = 384\n",
        "  config.hparams_mutation_ranges[\"ds_image_size\"] = [224, 384]\n",
        "  return config\n",
        "\n",
        "def get_config_vit_ti3():\n",
        "  config = get_common_config_vit()\n",
        "  config.batch_size = 512\n",
        "  config.vit_checkpoint_query = 'name==\"Ti/16\" and ds==\"i21k\" and aug==\"light1\" and wd==0.1 and sd==0.0'\n",
        "  config.hparams_defaults[\"num_layers\"] = 3\n",
        "  config.hparams_mutation_ranges[\"num_layers\"] = list(range(config.hparams_defaults[\"num_layers\"]+1))\n",
        "  config.hparams_defaults[\"ds_image_size\"] = 32\n",
        "  config.hparams_mutation_ranges[\"ds_image_size\"] = [16*i for i in (range(1, 1+int(112/16)))]\n",
        "  return config\n",
        "\n",
        "def get_config_vit_base():\n",
        "  config = get_common_config_vit()\n",
        "  config.batch_size = 256\n",
        "  config.vit_checkpoint_query = 'name==\"B/16\" and ds==\"i21k\" and aug==\"medium1\" and wd==0.1 and sd==0'\n",
        "  max_num_layers = get_max_num_layers(config.vit_checkpoint_query)\n",
        "  config.hparams_defaults[\"num_layers\"] = max_num_layers\n",
        "  config.hparams_mutation_ranges[\"num_layers\"] = list(range(config.hparams_defaults[\"num_layers\"]+1))\n",
        "  config.hparams_defaults[\"ds_image_size\"] = 80\n",
        "  config.hparams_mutation_ranges[\"ds_image_size\"] = [16*i for i in (range(1, 1+int(112/16)))]\n",
        "  return config\n",
        "\n",
        "def config_validate(config):\n",
        "  for khp in config.hparams_defaults:\n",
        "    if khp in config.hparams_mutation_ranges:\n",
        "      assert config.hparams_defaults[khp] in config.hparams_mutation_ranges[khp], khp\n",
        "  for khp in config.hparams_mutation_ranges:\n",
        "    assert khp in config.hparams_defaults, khp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_wzXhNikODR8"
      },
      "outputs": [],
      "source": [
        "class Vit(Agent):  # ViT large.\n",
        "  @staticmethod\n",
        "  def get_model_factory():\n",
        "    return VitModelFactory\n",
        "\n",
        "  def load_state(self):\n",
        "    task_name = self.config.task_name\n",
        "    self.pop = Population(self.config)\n",
        "    self.num_paths_sampled = 0\n",
        "    if task_name.startswith(\"root_model/\"):  # Root models.\n",
        "      hparams = self.config.hparams_defaults.as_configdict()\n",
        "      if task_name == \"root_model/random_init\":\n",
        "        hparams[\"num_classes\"] = 0  # Removes head layer.\n",
        "        path_params = self.get_model_factory().get_init_comps(hparams, self.config)\n",
        "      else:\n",
        "        assert task_name == \"root_model/checkpoint\", task_name\n",
        "        path_params = get_vit_checkpoint_comps(hparams[\"ds_image_size\"],\n",
        "                                               self.config.vit_checkpoint_query)\n",
        "      path = Path(hparams, [Component(k, self.id, v) for k, v in path_params.items()],\n",
        "                  parent=None, agent_id=self.id, task_name=task_name)\n",
        "      self.pop.paths[self.id].append(path)\n",
        "    else:\n",
        "      load_paths(self)\n",
        "      assert self.pop.paths, \"Empty population! run an agent creating a root model.\"\n",
        "      df_leaderboard(pop_to_df(self.pop))\n",
        "\n",
        "  def get_config(self):\n",
        "    return get_config_vit_large()\n",
        "\n",
        "  def run(self):\n",
        "    if self.config.task_name.startswith(\"root_model/\"):\n",
        "      run_root_model(self)\n",
        "      return\n",
        "    run_evolution(self)\n",
        "\n",
        "  def do_mutate(self, hparams, mutation_name):\n",
        "    \"\"\"Returns True if mutation (∂) is sampled to be applied.\"\"\"\n",
        "    if mutation_name in self.config.get(\"force_mutations\", []):\n",
        "      return True\n",
        "    mutation_prob_k = f\"_mu_|{mutation_name}\"   # Key for ∂ in µ() lookup table.\n",
        "    if \"_mu_\" in self.config.hparams_mutation_ranges:  # Mutate µ(∂) value.\n",
        "      if hparams[\"_mu_\"] \u003e np.random.uniform():\n",
        "        hparams[mutation_prob_k] = incremental_mutation(\n",
        "            hparams.get(mutation_prob_k, hparams[\"_mu_\"]),\n",
        "            self.config.hparams_mutation_ranges[\"_mu_\"])\n",
        "    return hparams.get(mutation_prob_k, hparams[\"_mu_\"]) \u003e np.random.uniform()\n",
        "\n",
        "  def parent_decay_selection(self):\n",
        "    parent = decay_selection(self.pop.paths[self.config.agent_id])\n",
        "    if parent is not None:\n",
        "      parent.metrics[\"offsprings\"] = parent.metrics.get(\"offsprings\", 0) + 1\n",
        "    return parent\n",
        "\n",
        "  def sample_path(self, gen_hparams):\n",
        "    parent = self.parent_decay_selection()\n",
        "    if not parent:  # Random sample.\n",
        "      parent = random.choice([p for paths in self.pop.paths.values() for p in paths])\n",
        "      print(f\"  Randomly selected parent {parent.agent_id}:{parent.id}\")\n",
        "    return self.mutate_parent(parent, gen_hparams)\n",
        "\n",
        "  def mutate_hparams(self, hparams):\n",
        "    for k in sorted(self.config.hparams_mutation_ranges):\n",
        "      if k in hparams and self.do_mutate(hparams, f\"hp:{k}\"):\n",
        "        hparams[k] = incremental_mutation(\n",
        "            hparams[k], self.config.hparams_mutation_ranges[k])\n",
        "\n",
        "  def is_generation_hp(self, k):  # True if the hp is shared across generation.\n",
        "    return (k.startswith(\"ds_\") or k.startswith(\"_mu_|hp:ds_\") or k in [\"_mu_\"])\n",
        "    \n",
        "  def sample_gen_hps(self):\n",
        "    \"\"\"Sample hparams that need to be shared across each paths generation.\"\"\"\n",
        "    gen_hparams = {k:v for k,v in self.config.hparams_defaults.items() if self.is_generation_hp(k)}\n",
        "    # Overwrite with values from best path if available.\n",
        "    best_path = self.pop.get_best_path()\n",
        "    if best_path:\n",
        "      gen_hparams.update({k:v for k,v in best_path.hparams.items() if self.is_generation_hp(k)})\n",
        "      self.mutate_hparams(gen_hparams)  # Sample mutations.\n",
        "    return gen_hparams\n",
        "\n",
        "  def finalize_hps(self, hparams, paths):\n",
        "    # Validate shared params.\n",
        "    for k in hparams:\n",
        "      if self.is_generation_hp(k):\n",
        "        for path in paths:\n",
        "          assert hparams[k] == path.hparams[k]\n",
        "    return hparams\n",
        "\n",
        "  def allow_component_reuse(self, comp_name, parent):\n",
        "    # Head must be trainable. If parent is not from same agent it will fall back\n",
        "    # to random init of correct shape.\n",
        "    return not(comp_name == \"head\" and self.config.agent_id != parent.agent_id)\n",
        "\n",
        "  def mutate_parent(self, parent, gen_hparams):\n",
        "    config = self.config\n",
        "    agent_id = config.agent_id\n",
        "    task_name = config.task_name\n",
        "    comps = []\n",
        "    new_hparams = copy.deepcopy(parent.hparams)\n",
        "    self.mutate_hparams(new_hparams)\n",
        "    # Overwrite hparams shared with paths in the current generation.\n",
        "    new_hparams.update(gen_hparams)\n",
        "\n",
        "    init_params = self.get_model_factory().get_init_comps(new_hparams, config)\n",
        "    # Attept to reuse matching componenent from parent.\n",
        "    comps_lookup = {c.name:c for c in parent.components}\n",
        "    for new_comp_name in init_params:\n",
        "      comp = None\n",
        "      if new_comp_name in comps_lookup and self.allow_component_reuse(new_comp_name, parent):\n",
        "        parent_comp = comps_lookup[new_comp_name]\n",
        "        # Check shapes match otherwise skip.\n",
        "        if (jax.tree_util.tree_map(jnp.shape, init_params[new_comp_name]) ==\n",
        "            jax.tree_util.tree_map(jnp.shape, parent_comp.params)):\n",
        "          if parent_comp.is_trainable() or self.do_mutate(new_hparams, f\"clone:{new_comp_name}\"):\n",
        "            comp = parent_comp.clone(agent_id)  # Clone trainable component.\n",
        "          else:\n",
        "            comp = parent_comp  # Refer to frozen component.\n",
        "        else:\n",
        "          if new_comp_name == \"posembed_input\":\n",
        "            # Change of image size changed shape of position embeddings,\n",
        "            # this can happend if ds_image_size is tuned.\n",
        "            assert \"ds_image_size\" in config.hparams_mutation_ranges\n",
        "            assert new_hparams[\"ds_image_size\"] != parent.hparams[\"ds_image_size\"]\n",
        "          else:\n",
        "            print(f\"WARNING: Shapes mismatch for component: {new_comp_name}  {parent.agent_id}-\u003e{agent_id}\")\n",
        "            print(jax.tree_util.tree_map(jnp.shape, init_params[new_comp_name]))\n",
        "            print(jax.tree_util.tree_map(jnp.shape, parent_comp.params))\n",
        "            assert False\n",
        "      if comp is None:  # Custom component mutation.\n",
        "        comp = self.mutate_component(new_comp_name, new_hparams, init_params, parent)\n",
        "      if comp is None:  # Random init.\n",
        "        print(f\"    Initialized parameters of \\\"{new_comp_name}\\\" module\")\n",
        "        comp = Component(new_comp_name, agent_id, init_params[new_comp_name])\n",
        "      comps.append(comp)\n",
        "    return Path(new_hparams, comps, parent=parent, agent_id=agent_id, task_name=task_name)\n",
        "\n",
        "  def mutate_component(self, comp_name, hparams, init_params, parent):\n",
        "    if comp_name == \"posembed_input\":\n",
        "      parent_posembed = [c for c in parent.components if c.name == comp_name]\n",
        "      assert len(parent_posembed) == 1\n",
        "      parent_posembed = parent_posembed[0].params[\"pos_embedding\"]\n",
        "      mapped_posembed = vit_checkpoint.interpolate_posembed(\n",
        "          parent_posembed, init_params[comp_name][\"pos_embedding\"].shape[1], True)\n",
        "      return Component(\n",
        "          name=\"posembed_input\", agent_id=self.config.agent_id,\n",
        "          params=flax.core.freeze({\"pos_embedding\": mapped_posembed}))\n",
        "\n",
        "  def print_best_path_summary(self):\n",
        "    best_path = self.pop.get_best_path()\n",
        "    print(f\"Best id:{best_path.id}\", f\"score:{best_path.score():.4f}\",\n",
        "          f\"quality:{best_path.metrics['quality']:.4f}\", f\"\\n{best_path.hparams}\")\n",
        "\n",
        "  def get_paths_to_publish(self):\n",
        "    return [self.pop.get_best_path()]\n",
        "\n",
        "class VitT3(Vit):  # ViT tiny 3 layers.\n",
        "  def get_config(self):\n",
        "    return get_config_vit_ti3()\n",
        "\n",
        "class VitB(Vit):  # ViT base.\n",
        "  def get_config(self):\n",
        "    return get_config_vit_base()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-N57gaotPN55"
      },
      "source": [
        "## MultiVit Agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cb3uj_kw9ED9"
      },
      "outputs": [],
      "source": [
        "def set_multivit_common_config(config):\n",
        "  config.force_mutations = []\n",
        "  config.scorer_kwargs = {}\n",
        "  for rm_k in [\"num_layers\", \"ds_image_size\"]:\n",
        "    del config.hparams_defaults[rm_k]\n",
        "    del config.hparams_mutation_ranges[rm_k]\n",
        "  config.hparams_defaults[\"router_init_main_path_weight\"] = 0.8\n",
        "  config.hparams_defaults[\"router_lr_mult\"] = 0.05\n",
        "  config.hparams_defaults[\"num_paths\"] = 2\n",
        "  config.hparams_mutation_ranges[\"router_lr_mult\"] = [0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1]\n",
        "  config.hparams_mutation_ranges[\"num_paths\"] = [2, 3]\n",
        "  return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Id10GxBcyJi"
      },
      "outputs": [],
      "source": [
        "class MultiVit(Vit):\n",
        "  def get_config(self):\n",
        "    config = get_config_vit_large()\n",
        "    return set_multivit_common_config(config)\n",
        "\n",
        "  @staticmethod\n",
        "  def get_model_factory():\n",
        "    return MultiVitModelFactory\n",
        "\n",
        "  def run(self):\n",
        "    run_evolution(self)\n",
        "\n",
        "  def agents_to_load(self):\n",
        "    return [self.single_path_agent_class()+\"~*\", self.config.agent_id]\n",
        "\n",
        "  def single_path_agent_class(self):\n",
        "    return self.class_name.removeprefix(\"Multi\")\n",
        "\n",
        "  def single_path_main_agent_id(self):\n",
        "    return format_agent_id(self.single_path_agent_class(), self.config.task_name)\n",
        "\n",
        "  def load_state(self):\n",
        "    if self.config.task_name.startswith(\"root_model/\"):\n",
        "      assert False, (\n",
        "          \"Root models need to be generated with the corresponding\" \\\n",
        "          f\"single path agent: '{self.single_path_agent_class()}'.\")\n",
        "    super().load_state()\n",
        "    assert self.single_path_main_agent_id() in self.pop.paths, (\n",
        "        \"Missing state for the corresponding single path main agent. \" \\\n",
        "        f\"Run agent '{self.single_path_agent_class()}' \" \\\n",
        "        f\"on the '{self.config.task_name}' task to generated it.\")\n",
        "\n",
        "  def sample_path(self, gen_hparams):\n",
        "    selected_paths = {}\n",
        "    selected_paths[\"path_0\"] = random.choice(self.pop.paths[self.single_path_main_agent_id()])\n",
        "    parent = self.parent_decay_selection()\n",
        "    if parent is not None:\n",
        "      new_hparams = copy.deepcopy(parent.hparams)\n",
        "    else:\n",
        "      parent = selected_paths[\"path_0\"]\n",
        "      best_path = self.pop.get_best_path()\n",
        "      if best_path:\n",
        "        new_hparams = copy.deepcopy(best_path.hparams)\n",
        "      else:\n",
        "        new_hparams = self.config.hparams_defaults.to_dict()\n",
        "      print(\"  New random\")\n",
        "    new_hparams[\"paths\"] = {}\n",
        "    self.mutate_hparams(new_hparams)\n",
        "    # Overwrite hparams shared with paths in the current generation.\n",
        "    new_hparams.update(gen_hparams)\n",
        "\n",
        "    for path_name in [f\"path_{i}\" for i in range(int(new_hparams[\"num_paths\"]))]:\n",
        "      if path_name in selected_paths:\n",
        "        continue\n",
        "      if path_name in parent.hparams.get(\"paths\", {}):\n",
        "        selected_paths[path_name] = self.pop.get_path_from_full_id(parent.hparams[\"paths\"][path_name][\"agent_id\"], parent.hparams[\"paths\"][path_name][\"id\"])\n",
        "        print(\"  Subpath from parent: \", path_name, selected_paths[path_name].full_id)\n",
        "        continue\n",
        "      selected_paths[path_name] = random.choice([\n",
        "          p for paths in self.pop.paths.values() for p in paths if (\n",
        "              p.agent_id not in [sp.agent_id for sp in selected_paths.values()]\n",
        "              and p.agent_id.startswith(f\"{self.single_path_agent_class()}~\")\n",
        "              and not p.agent_id.endswith(\"~1k\")  # Excludes VTAB-1k tasks.\n",
        "              # and (path_name != \"path_1\" or p.agent_id in ['Vit~i_naturalist2017'])  # Forces i_naturalist2017 selection.\n",
        "              )])\n",
        "      print(\"  Subpath rand selected:\", path_name, selected_paths[path_name].full_id)\n",
        "\n",
        "    for (path_name, path) in selected_paths.items():\n",
        "      new_hparams[\"paths\"][path_name] = {\n",
        "          \"id\": path.id,\n",
        "          \"agent_id\": path.agent_id,\n",
        "          \"hparams\": copy.deepcopy(path.hparams)}\n",
        "    print(\"  Sampled subpaths:\",\n",
        "          {k: v[\"agent_id\"] for k, v in new_hparams[\"paths\"].items()})\n",
        "    # Set headless model config for models of different tasks.\n",
        "    for k in new_hparams[\"paths\"]:\n",
        "      if selected_paths[k].task_name != self.config.task_name:\n",
        "        new_hparams[\"paths\"][k][\"hparams\"][\"num_classes\"] = 0\n",
        "    # Collect image sizes needed.\n",
        "    image_sizes = set()\n",
        "    for k in new_hparams[\"paths\"]:\n",
        "      image_sizes.add(int(new_hparams[\"paths\"][k][\"hparams\"][\"ds_image_size\"]))\n",
        "    new_hparams[\"ds_image_size\"] = list(image_sizes)\n",
        "    # Collect components.\n",
        "    init_params = self.get_model_factory().get_init_comps(new_hparams, self.config)\n",
        "    comps = []\n",
        "    for new_comp_name in init_params:\n",
        "      if new_comp_name in new_hparams[\"paths\"]:\n",
        "        comps.append(ComponentPath(name=new_comp_name,\n",
        "                                   path=selected_paths[new_comp_name]))\n",
        "      else:\n",
        "        comps_lookup = {c.name:c for c in parent.components}\n",
        "        if new_comp_name in comps_lookup and (\n",
        "          jax.tree_util.tree_map(jnp.shape, init_params[new_comp_name]) ==\n",
        "          jax.tree_util.tree_map(jnp.shape, comps_lookup[new_comp_name].params)):\n",
        "          print(\"  COMP Reusing\", new_comp_name)\n",
        "          comp = comps_lookup[new_comp_name].clone(agent_id=self.config.agent_id)\n",
        "        else:\n",
        "          print(\"  COMP Init\", new_comp_name)\n",
        "          comp = Component(name=new_comp_name, agent_id=self.config.agent_id,\n",
        "                           params=init_params[new_comp_name])\n",
        "        comps.append(comp)\n",
        "    return Path(new_hparams, comps, parent=parent,\n",
        "                agent_id=self.config.agent_id, task_name=self.config.task_name)\n",
        "\n",
        "  def finalize_hps(self, hparams, paths):\n",
        "    # Validate shared params.\n",
        "    for k in hparams:\n",
        "      if self.is_generation_hp(k) and k != \"ds_image_size\":\n",
        "        for path in paths:\n",
        "          assert hparams[k] == path.hparams[k], (k, hparams[k], path.hparams[k])\n",
        "    image_sizes = set()\n",
        "    for path in paths:\n",
        "      image_sizes.update(path.hparams[\"ds_image_size\"])\n",
        "    hparams[\"ds_image_size\"] = list(image_sizes)\n",
        "    print(\"Image sizes:\", hparams[\"ds_image_size\"])\n",
        "    return hparams\n",
        "\n",
        "  def print_best_path_summary(self):\n",
        "    super().print_best_path_summary()\n",
        "    best_path = self.pop.get_best_path()\n",
        "    print(\"Paths used by best model:\",\n",
        "          {k: best_path.hparams[\"paths\"][k][\"agent_id\"] for k in best_path.hparams[\"paths\"]})\n",
        "\n",
        "  def get_paths_to_publish(self):\n",
        "    paths_to_publish = [self.pop.get_best_path()]\n",
        "    for path in list(paths_to_publish):\n",
        "      for k in path.hparams[\"paths\"]:\n",
        "        paths_to_publish.append(self.pop.get_path_from_full_id(path.hparams[\"paths\"][k][\"agent_id\"], path.hparams[\"paths\"][k][\"id\"]))\n",
        "    return paths_to_publish\n",
        "\n",
        "class MultiVitT3(MultiVit):\n",
        "  def get_config(self):\n",
        "    return set_multivit_common_config(get_config_vit_ti3())\n",
        "\n",
        "class MultiVitB(MultiVit):\n",
        "  def get_config(self):\n",
        "    return set_multivit_common_config(get_config_vit_base())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gfGafBjhDpV7"
      },
      "source": [
        "# Tasks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mBAsZx4IWWaS"
      },
      "outputs": [],
      "source": [
        "TFDS_BUILDERS_CACHE = {}\n",
        "def get_tfds_builder(tfds_name):\n",
        "  global TFDS_BUILDERS_CACHE\n",
        "  if tfds_name not in TFDS_BUILDERS_CACHE:\n",
        "    TFDS_BUILDERS_CACHE[tfds_name] = tfds.builder(tfds_name)\n",
        "    TFDS_BUILDERS_CACHE[tfds_name].download_and_prepare()\n",
        "  return TFDS_BUILDERS_CACHE[tfds_name]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqnhOnk5jbnm"
      },
      "outputs": [],
      "source": [
        "def get_versioned_tfds_name(tfds_name):\n",
        "  version = TFDS_IMAGE_CLASSIFCATON_TASKS.get(tfds_name, {}).get(\"version\", \"\")\n",
        "  if version:\n",
        "    return tfds_name + \":\" + version\n",
        "  return tfds_name\n",
        "\n",
        "def get_default_splits(tfds_name):\n",
        "  info = get_tfds_builder(tfds_name).info\n",
        "  splits = list(info.splits.keys())\n",
        "  assert \"train\" in splits, splits\n",
        "  splits.remove(\"train\")\n",
        "  used_percent = 0\n",
        "  slice_percent = 5\n",
        "  pp = {}\n",
        "  for k in [\"test\", \"validation\"]:\n",
        "    if k in splits:\n",
        "      pp[k] = k\n",
        "      splits.remove(k)\n",
        "    else:\n",
        "      pp[k] = f\"train[{used_percent}%:{used_percent+slice_percent}%]\"\n",
        "      used_percent += slice_percent\n",
        "  pp[\"train\"] = f\"train[{used_percent}%:]\"\n",
        "  return pp\n",
        "\n",
        "def get_dataset_and_splits(tfds_name: str):\n",
        "  vtab_class = None\n",
        "  if tfds_name in [\"imagenet_v2\", \"cifar10_1\"]:\n",
        "    assert False,  f\"{tfds_name} used as validation set for other tasks.\"\n",
        "  dataset = get_versioned_tfds_name(tfds_name)\n",
        "  if tfds_name == \"imagenet2012\":\n",
        "    dataset = {\"train\":dataset, \"validation\":\"imagenet_v2:3.0.0\", \"test\":dataset}\n",
        "    splits = {\"train\":\"train\", \"validation\":\"test\", \"test\":\"validation\"}\n",
        "  elif tfds_name == \"cifar100\":\n",
        "    splits = {\"train\":\"train[:98%]\", \"validation\":\"train[98%:]\", \"test\":\"test\"}\n",
        "  elif tfds_name == \"cifar10\":\n",
        "    dataset = {\"train\":dataset, \"validation\":\"cifar10_1:1.1.0\", \"test\":dataset}\n",
        "    splits = {\"train\":\"train\", \"validation\":\"test\", \"test\":\"test\"}\n",
        "  elif (tfds_name.startswith(\"visual_domain_decathlon/\") or\n",
        "        tfds_name in [\"i_naturalist2017\", \"i_naturalist2018\", \"places365_small\"]):\n",
        "    # Test has no labels, split validation in half.\n",
        "    splits =  {\"train\":\"train\", \"validation\":\"validation[:50%]\", \"test\":\"validation[50%:]\"}\n",
        "  elif tfds_name.startswith(\"cmaterdb/\"):\n",
        "    # Increase size of validation set due to small dataset size.\n",
        "    splits =  {\"train\":\"train[20%:]\", \"validation\":\"train[:20%]\", \"test\":\"test\"}\n",
        "  elif tfds_name == \"omniglot\":\n",
        "    # Test has no labels, and missing validation, use additional splits.\n",
        "    splits = {\"train\":\"train\", \"validation\":\"small1\", \"test\":\"small2\"}\n",
        "  elif tfds_name.startswith(\"controlled_noisy_web_labels/\"):\n",
        "    splits =  {\"train\":\"train_00\", \"validation\":\"validation[:50%]\", \"test\":\"validation[50%:]\"}\n",
        "  elif tfds_name.startswith(\"cycle_gan/\"):\n",
        "    splits =  {\"train\":\"trainA[10%:]+trainB[10%:]\",\n",
        "               \"validation\":\"trainA[:10%]+trainB[:10%]\",\n",
        "               \"test\":\"testA+testB\"}\n",
        "  elif tfds_name in [\"imagenet_a\", \"imagenet_r\", \"imagenet_sketch\",\n",
        "                     \"siscore/rotation\", \"siscore/size\", \"siscore/location\",]:\n",
        "    # Only test split.\n",
        "    splits =  {\"train\":\"test[10%:]\", \"validation\":\"test[5%:10%]\", \"test\":\"test[:5%]\"}\n",
        "  elif tfds_name in [\"pet_finder\"]:\n",
        "    # Explicitly use only train split. E.g. test has no labels.\n",
        "    splits =  {\n",
        "        \"train\":\"train[10%:]\",\n",
        "        \"validation\":\"train[5%:10%]\",\n",
        "        \"test\":\"train[:5%]\"}\n",
        "  elif tfds_name == \"quickdraw_bitmap\":\n",
        "    # Cap size of test and validation set.\n",
        "    splits =  {\"train\":\"train[20000:]\", \"validation\":\"train[10000:20000]\", \"test\":\"train[:10000]\"}\n",
        "  elif tfds_name == \"stanford_online_products\":\n",
        "    # Use the first 10k test samples as validation since test has 60k.\n",
        "    splits =  {\"train\":\"train\", \"validation\":\"test[:10000]\", \"test\":\"test[10000:]\"}\n",
        "  elif tfds_name in VTAB_TASKS or (\n",
        "      tfds_name.endswith(\"/1k\") and tfds_name.replace(\"/1k\", \"\") in VTAB_TASKS):\n",
        "    is_vtab_1k = tfds_name.endswith(\"/1k\")\n",
        "    tfds_name = tfds_name.replace(\"/1k\", \"\")\n",
        "    registry_name = {\n",
        "        \"diabetic_retinopathy_detection/btgraham-300\": \"diabetic_retinopathy\",\n",
        "        \"svhn_cropped\": \"svhn\",\n",
        "        \"cifar100\": \"cifar\",\n",
        "        \"cifar10\": \"cifar\",\n",
        "    }.get(tfds_name, tfds_name.split(\"/\")[0])\n",
        "    args = {\n",
        "        \"clevr/count_all\": (\"count_all\",),\n",
        "        \"clevr/count_cylinders\": (\"count_cylinders\",),\n",
        "        \"clevr/closest_object_distance\": (\"closest_object_distance\",),\n",
        "        \"dsprites/label_x_position\": (\"label_x_position\",),\n",
        "        \"dsprites/label_orientation\": (\"label_orientation\",),\n",
        "        \"kitti/closest_object_distance\": (\"closest_object_distance\",),\n",
        "        \"kitti/count_vehicles\": (\"count_vehicles\",),\n",
        "        \"kitti/closest_vehicle_distance\": (\"closest_vehicle_distance\",),\n",
        "        \"smallnorb/label_category\": (\"label_category\",),\n",
        "        \"smallnorb/label_lighting\": (\"label_lighting\",),\n",
        "        \"smallnorb/label_azimuth\": (\"label_azimuth\",),\n",
        "        \"smallnorb/label_elevation\": (\"label_elevation\",),\n",
        "        \"cifar100\": (100,),\n",
        "        \"cifar10\": (10,),\n",
        "    }.get(tfds_name, ())\n",
        "    vtab_class = task_adapt_registry.Registry.lookup(\n",
        "        f\"data.{registry_name}\")(*args)\n",
        "    vtab_splits = vtab_class._tfds_splits\n",
        "    dataset = {\n",
        "        \"caltech101\": \"caltech101:3.*.*\",\n",
        "        \"dtd\": \"dtd:3.*.*\",\n",
        "        \"oxford_flowers102\": \"oxford_flowers102:2.*.*\",\n",
        "        \"oxford_iiit_pet\": \"oxford_iiit_pet:3.*.*\",\n",
        "        \"sun397\": \"sun397/tfds:4.*.*\",\n",
        "        \"svhn\": \"svhn_cropped:3.*.*\",\n",
        "        \"patch_camelyon\": \"patch_camelyon:2.*.*\",\n",
        "        \"eurosat\": \"eurosat/rgb:2.*.*\",\n",
        "        \"resisc45\": \"resisc45:3.*.*\",\n",
        "        \"diabetic_retinopathy\": \"diabetic_retinopathy_detection/btgraham-300:3.*.*\",\n",
        "        \"clevr\": \"clevr:3.*.*\",\n",
        "        \"dmlab\": \"dmlab:2.0.1\",\n",
        "        \"dsprites\": \"dsprites:2.*.*\",\n",
        "        \"kitti\": \"kitti:3.2.0\",\n",
        "        \"smallnorb\": \"smallnorb:2.*.*\",\n",
        "        \"cifar\" : \"cifar100:3.*.*\" if tfds_name == \"cifar100\" else \"cifar10:3.*.*\",\n",
        "    }[registry_name]\n",
        "    if is_vtab_1k:\n",
        "      splits =  {\"train\": str(vtab_splits[\"train800\"]),\n",
        "                 \"validation\": str(vtab_splits[\"val200\"]),\n",
        "                 \"test\": str(vtab_splits[\"test\"])}\n",
        "    else:\n",
        "      splits =  {\"train\": str(vtab_splits[\"train\"]),\n",
        "                 \"validation\": str(vtab_splits[\"val\"]),\n",
        "                 \"test\": str(vtab_splits[\"test\"])}\n",
        "  else:\n",
        "    splits = get_default_splits(tfds_name)\n",
        "  return dataset, splits, vtab_class\n",
        "\n",
        "class Task():\n",
        "  def __init__(self, name, config):\n",
        "    self.config = config\n",
        "\n",
        "    self.dataset, self.splits, self.vtab_class = get_dataset_and_splits(name)\n",
        "    self.name = name\n",
        "    if self.vtab_class:\n",
        "      self.num_classes = self.vtab_class.get_num_classes()\n",
        "    else:\n",
        "      self.num_classes = self.get_builder(\n",
        "          \"train\").info.features[self.get_label_key()].num_classes\n",
        "    num_train_examples = self.get_builder(\n",
        "        \"train\").info.splits[self.splits[\"train\"]].num_examples\n",
        "    self.train_batch_size = config.batch_size\n",
        "    self.num_train_batches_between_validations = math.ceil(\n",
        "        min(num_train_examples,\n",
        "            config.num_train_examples_between_validations_max)\n",
        "        / self.train_batch_size)\n",
        "\n",
        "    num_validation_examples_tot = self.get_builder(\n",
        "        \"validation\").info.splits[self.splits[\"validation\"]].num_examples\n",
        "    if config.num_validation_examples_max \u003c= num_validation_examples_tot:\n",
        "      self.validation_batch_size = config.batch_size\n",
        "      self.num_validation_batches = math.floor(\n",
        "          config.num_validation_examples_max / self.validation_batch_size)\n",
        "    else:\n",
        "      # Adjust batch_size and num_batches to cover the smaller validation sets.\n",
        "      self.num_validation_batches = math.ceil(\n",
        "          num_validation_examples_tot / config.batch_size)\n",
        "      self.validation_batch_size = math.floor(\n",
        "          num_validation_examples_tot / self.num_validation_batches)\n",
        "      assert num_validation_examples_tot \u003e= (\n",
        "          self.num_validation_batches*self.validation_batch_size)\n",
        "    self.num_validation_examples = (\n",
        "        self.num_validation_batches * self.validation_batch_size)\n",
        "\n",
        "    print(f\"Task: {self.name}\")\n",
        "    print(f\"  Train batches between validations: {self.num_train_batches_between_validations}\")\n",
        "    print(f\"  Validation batches: {self.num_validation_batches}\")\n",
        "    print(f\"  Validation batch size: {self.validation_batch_size}\")\n",
        "    print(f\"  Dataset {self.dataset}\")\n",
        "    print(f\"  Splits {self.splits}\")\n",
        "\n",
        "  def get_label_key(self):\n",
        "    return {\n",
        "        \"stanford_online_products\": \"super_class_id\",\n",
        "        }.get(self.name, \"label\")\n",
        "\n",
        "  def get_builder(self, mode):\n",
        "    if type(self.dataset) == str:\n",
        "      return get_tfds_builder(self.dataset)\n",
        "    return get_tfds_builder(self.dataset[mode])\n",
        "\n",
        "  def get_ds(self, mode, hparams):\n",
        "    data = self.get_builder(mode).as_dataset(\n",
        "        split=self.splits[mode],\n",
        "        shuffle_files=mode==\"train\")\n",
        "\n",
        "    def _pp(data):  # Preprocessing function.\n",
        "      im = data[\"image\"]\n",
        "      tf.debugging.assert_type(im, tf.uint8)\n",
        "\n",
        "      if mode == \"train\":\n",
        "        if hparams.get(\"ds_quality_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_jpeg_quality(\n",
        "              im,\n",
        "              min_jpeg_quality=int(100 * (1 - hparams[\"ds_quality_delta\"])),\n",
        "              max_jpeg_quality=100)\n",
        "\n",
        "      # Must have 3 channels.\n",
        "      if im.shape[-1] == 1:\n",
        "        im = tf.squeeze(tf.stack([im] * 3, -1), axis=-2)\n",
        "      assert im.shape[-1] == 3\n",
        "      im = tf.cast(im, tf.float32)\n",
        "      if mode == \"train\":\n",
        "        if hparams.get(\"ds_area_range_min\", 1.0) \u003c 1.0:\n",
        "          begin, size, _ = tf.image.sample_distorted_bounding_box(\n",
        "              tf.shape(im),\n",
        "              tf.zeros([0, 0, 4], tf.float32),\n",
        "              aspect_ratio_range=[hparams[\"ds_aspect_ratio_range_min\"],\n",
        "                                  1.0/hparams[\"ds_aspect_ratio_range_min\"]],\n",
        "              area_range=[hparams[\"ds_area_range_min\"], 1.0],\n",
        "              min_object_covered=0,  # Min image overlap already set by area_range.\n",
        "              use_image_if_no_bounding_boxes=True)\n",
        "          im = tf.slice(im, begin, size)\n",
        "          im.set_shape([None, None, 3]) # Restore dimension lost by tf.slice().\n",
        "        if hparams.get(\"ds_flip_left_right\", False):\n",
        "          if tf.random.uniform(shape=[]) \u003e 0.5:\n",
        "            im = tf.image.flip_left_right(im)\n",
        "        if hparams.get(\"ds_brightness_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_brightness(\n",
        "              im, max_delta=hparams[\"ds_brightness_delta\"])\n",
        "        if hparams.get(\"ds_contrast_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_contrast(\n",
        "              im, lower=1-hparams[\"ds_contrast_delta\"],\n",
        "              upper=1+hparams[\"ds_contrast_delta\"])\n",
        "        if hparams.get(\"ds_saturation_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_saturation(\n",
        "              im, lower=1-hparams[\"ds_saturation_delta\"],\n",
        "              upper=1+hparams[\"ds_saturation_delta\"])\n",
        "        if hparams.get(\"ds_hue_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_hue(im, max_delta=hparams[\"ds_hue_delta\"])\n",
        "\n",
        "      def get_formatted_image(image, image_size):\n",
        "        image = tf.image.resize(image, [image_size, image_size])\n",
        "        # Values in range [-1 , 1].\n",
        "        image = image / 127.5 - 1\n",
        "        image = tf.clip_by_value(image, -1, 1)\n",
        "        return image\n",
        "\n",
        "      if type(hparams[\"ds_image_size\"]) is list:\n",
        "        out_im = {}\n",
        "        for im_size in hparams[\"ds_image_size\"]:\n",
        "          out_im[str(im_size)] = get_formatted_image(im, int(im_size))\n",
        "      else:\n",
        "        out_im = get_formatted_image(im, int(hparams[\"ds_image_size\"]))\n",
        "      return {\"image\": out_im,\n",
        "              \"label\": data[self.get_label_key()]}\n",
        "\n",
        "    if mode == \"validation\":\n",
        "      data = data.take(self.num_validation_examples).cache()\n",
        "    if mode != \"test\":\n",
        "      data = data.repeat()\n",
        "    if self.vtab_class and self.vtab_class._base_preprocess_fn:\n",
        "      data = data.map(self.vtab_class._base_preprocess_fn, tf.data.AUTOTUNE)\n",
        "    data = data.map(_pp, tf.data.AUTOTUNE)\n",
        "    if mode == \"train\":\n",
        "      batch_size = self.train_batch_size\n",
        "    else:\n",
        "      batch_size = self.validation_batch_size\n",
        "    data = data.batch(batch_size)\n",
        "    if mode == \"train\":\n",
        "      data = data.shuffle(10)\n",
        "    return tfds.as_numpy(data.prefetch(tf.data.AUTOTUNE))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GdzUzvhrWLer"
      },
      "source": [
        "# Components"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DNxjDX13_dm_"
      },
      "outputs": [],
      "source": [
        "class Component():\n",
        "  def reset_globals():\n",
        "    Component.counter = 0\n",
        "    Component.last_saved = -1\n",
        "\n",
        "  def __init__(self, name, agent_id, params, id=None):\n",
        "    self._name = name\n",
        "    self.agent_id = agent_id\n",
        "    self._params = jax.device_get(params)\n",
        "    self._num_params = None\n",
        "    if id is None:\n",
        "      self.id = Component.counter\n",
        "      Component.counter += 1\n",
        "    else:\n",
        "      self.id = id\n",
        "    self.full_id = f\"{self.agent_id}:{self.id}\"\n",
        "\n",
        "  @property\n",
        "  def name(self):\n",
        "    if self._name is None:\n",
        "      self._name = get_comp_data(self, \"name\")\n",
        "    return self._name\n",
        "\n",
        "  @property\n",
        "  def params(self):\n",
        "    if self._params is None:\n",
        "      self._params = get_comp_params(self)\n",
        "    return self._params\n",
        "\n",
        "  @property\n",
        "  def num_params(self):\n",
        "    if self._num_params is None:\n",
        "      self._num_params = get_num_params(self.params)\n",
        "    return self._num_params\n",
        "\n",
        "  def is_trainable(self):\n",
        "    return len(Population.unmutable_comp2agents[self.full_id]) == 0\n",
        "\n",
        "  def clone(self, agent_id):\n",
        "    return Component(name=self.name, agent_id=agent_id,\n",
        "                     params=copy.deepcopy(jax.device_get(self.params)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SdMf8WcYBbjt"
      },
      "outputs": [],
      "source": [
        "class ComponentPath():\n",
        "  \"\"\"Wraps a Paths to be used as a Component.\"\"\"\n",
        "  def __init__(self, name, path):\n",
        "    self.name = name\n",
        "    self.path = path\n",
        "\n",
        "  def is_trainable(self):\n",
        "    return False\n",
        "\n",
        "  @property\n",
        "  def params(self):\n",
        "    return flax.core.freeze(self.path.get_all_params())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u9JsNREOCcpF"
      },
      "outputs": [],
      "source": [
        "AID2CDF = {}  # Maps agent ids to a dataframe with components metadata.\n",
        "\n",
        "def load_sharded_df(agent_dir, filename = \"components\"):\n",
        "  dfs = []\n",
        "  threads = []\n",
        "  lock = Lock()\n",
        "  def load_df(s_dir):\n",
        "    df = df_read_from_csv(s_dir, filename)\n",
        "    with lock:\n",
        "      dfs.append(df)\n",
        "  for s_dir in gfile.glob(agent_dir + f\"/state_*\"):\n",
        "    threads.append(Thread(target=load_df, args=(s_dir,)))\n",
        "    threads[-1].start()\n",
        "  [t.join() for t in threads]\n",
        "  return pd.concat(dfs)\n",
        "\n",
        "def get_comp_data(comp, data_name):\n",
        "  if comp.agent_id not in AID2CDF or comp.id not in AID2CDF[comp.agent_id]['id'].values:\n",
        "    AID2CDF[comp.agent_id] = load_sharded_df(os.path.join(SYSTEM_STATE_DIR, comp.agent_id))\n",
        "  c_data = AID2CDF[comp.agent_id].loc[AID2CDF[comp.agent_id]['id'] == comp.id]\n",
        "  assert len(c_data) == 1\n",
        "  return c_data[data_name].values[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0S5Baq2pQFSU"
      },
      "outputs": [],
      "source": [
        "CFID2PARAMS = {}  # Comps full id 2 comp parameters.\n",
        "\n",
        "def get_comp_params(comp):\n",
        "  if comp.full_id not in CFID2PARAMS:\n",
        "    stt = time.time()\n",
        "    state_iid = get_comp_data(comp, \"state_iid\")\n",
        "    agent_dir = os.path.join(SYSTEM_STATE_DIR, comp.agent_id)\n",
        "    chkpnt = os.path.join(agent_dir, f\"state_{state_iid}/checkpoint_{state_iid}\")\n",
        "    loaded_params = flax.core.freeze(flax_checkpoints.restore_checkpoint(\n",
        "        ckpt_dir=os.path.dirname(chkpnt), target=None))\n",
        "    for k in loaded_params.keys():\n",
        "      _, c_agent_id, c_id = k.split(\":\")\n",
        "      full_id = f\"{c_agent_id}:{c_id}\"\n",
        "      assert full_id not in CFID2PARAMS, full_id\n",
        "      CFID2PARAMS[full_id] = loaded_params[k]\n",
        "    print(f\"  Loaded components: {time.time()-stt:.2f}s\", chkpnt, loaded_params.keys())\n",
        "  return CFID2PARAMS[comp.full_id]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pCie1hQLUyP7"
      },
      "source": [
        " # Paths \u0026 Population"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8PU5ffvd_gC9"
      },
      "outputs": [],
      "source": [
        "class Path():\n",
        "  def reset_globals(config):\n",
        "    Path.config = config\n",
        "    Path.counter = 0\n",
        "    Path.last_saved = -1\n",
        "    Path.paths = []\n",
        "    Path.scorer = ScorerDecay(**config.get(\"scorer_kwargs\", {}))\n",
        "\n",
        "  def __init__(self, hparams, components, parent, agent_id, task_name, id=None):\n",
        "    self.components = components\n",
        "    if id is None:\n",
        "      self.id = Path.counter\n",
        "      Path.counter += 1\n",
        "    else:\n",
        "      self.id = id\n",
        "    self.agent_id = agent_id\n",
        "    self.full_id = f\"{self.agent_id}:{self.id}\"\n",
        "    self.task_name = task_name\n",
        "    self.hparams = hparams\n",
        "    self.parent = parent\n",
        "    if self.parent is not None:\n",
        "      self.hparams[\"parent_agent_id\"] = parent.agent_id\n",
        "      self.hparams[\"parent_id\"] = int(parent.id)\n",
        "    self._model = None\n",
        "    self.metrics = {\"generation\": parent.metrics[\"generation\"]+1 if parent else 0}\n",
        "    Path.paths.append(self)\n",
        "\n",
        "  @property\n",
        "  def model_factory(self):\n",
        "    return get_agent_class(self.agent_id).get_model_factory()\n",
        "\n",
        "  @property\n",
        "  def model(self):\n",
        "    if self._model == None:\n",
        "      self._model = self.model_factory.get_model(self.hparams, self.config)\n",
        "    return self._model\n",
        "\n",
        "  def comps_only(self):  # Exclude wrapped paths.\n",
        "    return [c for c in self.components if c.__class__ is Component]\n",
        "\n",
        "  def score(self):\n",
        "    return Path.scorer.score(self)\n",
        "\n",
        "  def get_all_params(self):\n",
        "    params = {}\n",
        "    for c in self.components:\n",
        "      assert c.name not in params, c.name\n",
        "      params[c.name] = c.params\n",
        "    return flax.core.freeze(params)\n",
        "\n",
        "  def get_trainable_params(self):\n",
        "    params = {}\n",
        "    for c in self.components:\n",
        "      if c.is_trainable():\n",
        "        assert c.name not in params, c.name\n",
        "        params[c.name] = c.params\n",
        "    return flax.core.freeze(params)\n",
        "\n",
        "  def get_fixed_params(self):\n",
        "    params = {}\n",
        "    for c in self.components:\n",
        "      if not c.is_trainable():\n",
        "        assert c.name not in params, c.name\n",
        "        params[c.name] = c.params\n",
        "    return flax.core.freeze(params)\n",
        "\n",
        "  def update_trainable(self, trained_params):\n",
        "    trainable_count = 0\n",
        "    for c in self.components:\n",
        "      if c.is_trainable():\n",
        "        trainable_count += 1\n",
        "        assert c.name in trained_params.keys()\n",
        "        c._params = trained_params[c.name]\n",
        "    assert len(trained_params.keys()) == trainable_count, (\n",
        "        f\"{len(trained_params.keys())} {trainable_count}\")\n",
        "\n",
        "  def get_num_accounted_params(self):\n",
        "    rtn = 0\n",
        "    for c in sorted(self.components, key=lambda c: c.id):\n",
        "      agnts = copy.copy(Population.unmutable_comp2agents[c.full_id])\n",
        "      agnts.add(self.agent_id)\n",
        "      rtn += c.num_params/len(agnts)\n",
        "    return rtn\n",
        "\n",
        "  def get_flops(self):\n",
        "    return compute_flops_hlo(\n",
        "          partial(self.model.apply, train=False),\n",
        "          {\"params\": self.model_factory.get_comps2model_fn()(merge_params(\n",
        "              self.get_trainable_params(),\n",
        "              self.get_fixed_params()))},\n",
        "          self.model_factory.get_sample_input(self.hparams))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bkRmcJgzUbwN"
      },
      "outputs": [],
      "source": [
        "class ScorerDecay():\n",
        "  def __init__(self, scale_factor=1, base_accounted_params=0, base_flops=0):\n",
        "    assert 0.0 \u003c scale_factor \u003c= 1.0\n",
        "    self.scale_factor = scale_factor\n",
        "    self.base_accounted_params = base_accounted_params\n",
        "    self.base_flops = base_flops\n",
        "\n",
        "  def score(self, path):\n",
        "    if \"quality\" not in path.metrics or math.isnan(path.metrics[\"quality\"]):\n",
        "      return None\n",
        "    assert path.metrics[\"quality\"] \u003e= 0, (\n",
        "        f\"{path.task_name} {path.metrics['quality']}\")\n",
        "    score = path.metrics[\"quality\"]\n",
        "    if self.base_accounted_params \u003e 0:\n",
        "      if \"accounted_params\" not in path.metrics:\n",
        "        path.metrics[\"accounted_params\"] = path.get_num_accounted_params()\n",
        "      score *= self.scale_factor ** (path.metrics[\"accounted_params\"]/self.base_accounted_params)\n",
        "    if self.base_flops \u003e 0:\n",
        "      if \"flops\" not in path.metrics:\n",
        "        path.metrics[\"flops\"] = path.get_flops()\n",
        "      score *= self.scale_factor ** (path.metrics[\"flops\"]/self.base_flops)\n",
        "    assert score \u003e= 0\n",
        "    path.metrics[\"score\"] = score\n",
        "    return score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ePgFGuWPDts-"
      },
      "outputs": [],
      "source": [
        "def get_best_path(paths):\n",
        "  if paths:  # Oldest path achieving max score.\n",
        "    return max(sorted(paths, key=lambda p: p.id), key=lambda p: p.score())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YMbYgKd8_nyi"
      },
      "outputs": [],
      "source": [
        "class Population():\n",
        "  def __init__(self, config):\n",
        "    self.config = config\n",
        "    Path.reset_globals(config)\n",
        "    Component.reset_globals()\n",
        "    self.paths = defaultdict(list)\n",
        "    self.unmutable_comp2agents = None\n",
        "\n",
        "  def get_best_path(self):\n",
        "    return get_best_path(self.paths[self.config.agent_id])\n",
        "\n",
        "  def get_path_from_full_id(self, agent_id, path_id):\n",
        "    for p in self.paths[agent_id]:\n",
        "      if p.id == path_id:\n",
        "        return p\n",
        "    assert False, f\"Path not found {agent_id}:{path_id}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4EgQHbawpNcS"
      },
      "outputs": [],
      "source": [
        "pd.set_option(\"display.expand_frame_repr\", False)\n",
        "pd.set_option(\"display.max_columns\", 100)\n",
        "\n",
        "def pop_to_df(pop):\n",
        "  return paths_to_df([p for paths in pop.paths.values() for p in paths])\n",
        "\n",
        "def paths_to_df(paths):\n",
        "  def _format(x):\n",
        "    if type(x) in [dict, list]:\n",
        "      return json.dumps(x)\n",
        "    return x\n",
        "  metrics_keys = set()\n",
        "  hparams_keys = set()\n",
        "  for path in paths:\n",
        "    metrics_keys.update(path.metrics)\n",
        "    hparams_keys.update(path.hparams)\n",
        "  data = defaultdict(list)\n",
        "  for path in paths:\n",
        "    data[\"agent_id\"].append(path.agent_id)\n",
        "    data[\"task_name\"].append(path.task_name)\n",
        "    data[\"id\"].append(path.id)\n",
        "    data[\"components\"].append(\",\".join(\n",
        "        [f\"{c.agent_id}:{c.id}\" for c in path.comps_only()]))\n",
        "    for k in hparams_keys:\n",
        "      data[f\"hparams.{k}\"].append(_format(path.hparams[k]) if k in path.hparams else None)\n",
        "    for k in metrics_keys:\n",
        "      data[f\"metrics.{k}\"].append(path.metrics[k] if k in path.metrics else None)\n",
        "  return pd.DataFrame(data)\n",
        "\n",
        "def components_to_df(paths, agent_id, last_saved_iid, state_iid):\n",
        "  comps = set()\n",
        "  for p in paths:\n",
        "    comps.update(p.comps_only())\n",
        "  data = defaultdict(list)\n",
        "  for c in comps:\n",
        "    if c.agent_id != agent_id or c.id \u003c= last_saved_iid:\n",
        "      continue\n",
        "    data[\"agent_id\"].append(c.agent_id)\n",
        "    data[\"id\"].append(c.id)\n",
        "    data[\"name\"].append(c.name)\n",
        "    data[\"num_params\"].append(c.num_params)\n",
        "    data[\"state_iid\"].append(state_iid)\n",
        "  return pd.DataFrame(data)\n",
        "\n",
        "def print_df_segments(df, segment_length = 5):\n",
        "  tot_length = df.shape[0]\n",
        "  # Pad column title with spaces to keep alignment across segments.\n",
        "  def prepend_spaces(original_str, pad_to_len):\n",
        "    return \" \" * (pad_to_len-len(original_str)) + original_str\n",
        "  pad_to_len = max([len(tn) for tn in set(df[\"agent_id\"].to_list())])+1\n",
        "  df = df.rename(columns={\"agent_id\": prepend_spaces(\"agent_id\", pad_to_len),\n",
        "                          \"task_name\": prepend_spaces(\"task_name\", pad_to_len),\n",
        "                          \"hparams.parent_agent_id\": prepend_spaces(\"hparams.parent_agent_id\", pad_to_len)})\n",
        "  for x in range(0, tot_length, segment_length):\n",
        "    print(df[x:min(x+segment_length, tot_length)])\n",
        "\n",
        "def df_leaderboard(df):\n",
        "  # Place columns on the left for readability.\n",
        "  all_keys = sorted(df.columns.tolist())\n",
        "  first_keys = [\"agent_id\", \"id\", \"task_name\", \"metrics.test_quality\", \"metrics.score\",\n",
        "                \"metrics.quality\", \"metrics.accounted_params\", \"metrics.flops\",\n",
        "                \"hparams.parent_agent_id\", \"hparams.parent_id\"]\n",
        "  first_keys = [k for k in first_keys if k in all_keys]\n",
        "  sorted_keys = first_keys + [k for k in all_keys if k not in first_keys]\n",
        "  sorted_keys = [k for k in sorted_keys if \"_mu_|\" not in k]  # Remove mu function parameters.\n",
        "  df = df[sorted_keys]\n",
        "  if \"metrics.score\" in df:\n",
        "    df = df.sort_values([\"agent_id\", \"metrics.score\"], ascending=[True, False], ignore_index=True)\n",
        "  else:\n",
        "    df = df.sort_values(\"agent_id\", ignore_index=True)\n",
        "  print_df_segments(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rDU0butVlu2f"
      },
      "source": [
        "# Checkpointing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "104ReuGLspoJ"
      },
      "outputs": [],
      "source": [
        "def df_write_to_csv(df, dir_path, df_name):\n",
        "  filename_df = os.path.join(dir_path, f\"{df_name}.csv\")\n",
        "  with gfile.GFile(filename_df, \"w\") as outfile:\n",
        "    df.to_csv(outfile, index=False)\n",
        "\n",
        "def df_read_from_csv(dir_path, df_name):\n",
        "  filename_df = os.path.join(dir_path, f\"{df_name}.csv\")\n",
        "  with gfile.GFile(filename_df, \"r\") as infile:\n",
        "    df = pd.read_csv(infile)\n",
        "  # Pandas read_csv() reads empty stings as NaNs. Set NaNs to empty strings in\n",
        "  # columns with type strings/object.\n",
        "  for c in df.columns:\n",
        "    if df[c].dtype == np.object_:\n",
        "        df[c].fillna(\"\", inplace=True)\n",
        "  return df\n",
        "\n",
        "def get_comps_params_to_save(pop):\n",
        "  comps_params = {}\n",
        "  # All components generated by this agent.\n",
        "  all_comps = set(\n",
        "      [c for p in pop.paths[pop.config.agent_id] for c in p.comps_only() if c.agent_id == pop.config.agent_id])\n",
        "  # Check that there are not duplicate ids.\n",
        "  assert len(all_comps) == len(set([c.id for c in all_comps])), (\n",
        "      sorted([f\"{c.name}:{c.agent_id}:{c.id}\" for c in all_comps], key=lambda c: c.id))\n",
        "  for c in all_comps:\n",
        "    if c.id \u003c= Component.last_saved:\n",
        "      continue\n",
        "    assert c.agent_id == pop.config.agent_id\n",
        "    c_id_string = f\"{c.name}:{c.agent_id}:{c.id}\"\n",
        "    comps_params[c_id_string] = c.params\n",
        "  return comps_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z_AmkV7hLeQp"
      },
      "outputs": [],
      "source": [
        "def latest_checkpoint(ckpt_dir, prefix = \"checkpoint_\"):\n",
        "  ckpt_dir = os.fspath(ckpt_dir)\n",
        "  glob_path = os.path.join(ckpt_dir, f\"{prefix}*\")\n",
        "  checkpoint_files = flax_checkpoints.natural_sort(gfile.glob(glob_path))\n",
        "  checkpoint_files = [f for f in checkpoint_files if not f.endswith(\"_tmp\")]\n",
        "  return checkpoint_files[-1] if checkpoint_files else None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M_RTFjoOwfpM"
      },
      "outputs": [],
      "source": [
        "def save_checkpoint(ckpt_dir, comps_params, num_paths_sampled):\n",
        "  print(\"  Saving components\", num_paths_sampled, comps_params.keys())\n",
        "  flax_checkpoints.save_checkpoint(\n",
        "      ckpt_dir, target=comps_params, step=num_paths_sampled, overwrite=True)\n",
        "\n",
        "def save_state(agent):\n",
        "  write_start = time.time()\n",
        "  # Save data needed to resume exp.\n",
        "  state_dir = os.path.join(agent.config.agent_dir, f\"state_{agent.num_paths_sampled}\")\n",
        "  gfile.makedirs(state_dir)\n",
        "  assert not latest_checkpoint(state_dir), f\"Checkpoint already present in forlder: {state_dir}\"\n",
        "  print(\"Saving state:\", agent.num_paths_sampled)\n",
        "  df_write_to_csv(paths_to_df(agent.get_paths_to_publish()), state_dir, \"published\")\n",
        "  df_write_to_csv(paths_to_df([p for p in agent.pop.paths[agent.config.agent_id]]), state_dir, \"population\")\n",
        "  df_write_to_csv(paths_to_df(Path.paths).query(f'agent_id==\"{agent.config.agent_id}\" and id\u003e{Path.last_saved}'), state_dir, \"paths\")\n",
        "  df_write_to_csv(components_to_df(Path.paths, agent.config.agent_id, Component.last_saved, agent.num_paths_sampled), state_dir, \"components\")\n",
        "  json.dump(agent.config.as_configdict().to_dict(), gfile.GFile(os.path.join(state_dir, \"config.json\"), \"w\"), indent=2)\n",
        "  save_checkpoint(state_dir, get_comps_params_to_save(agent.pop), agent.num_paths_sampled)\n",
        "  print(f\"  State save time: {time.time() - write_start:.2f}s\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7yWdy7DskBph"
      },
      "outputs": [],
      "source": [
        "def load_paths(agent):\n",
        "  # Load system state metadata.\n",
        "  population_df = []\n",
        "  Population.unmutable_comp2agents = defaultdict(set)\n",
        "  for agent_to_load in agent.agents_to_load():\n",
        "    for agent_dir in gfile.glob(os.path.join(agent.config.system_state_dir, agent_to_load)):\n",
        "      agent_checkpoint = latest_checkpoint(os.path.join(agent_dir, \"state_*/\"))\n",
        "      if agent_checkpoint:\n",
        "        add_pdf = df_read_from_csv(os.path.dirname(agent_checkpoint), \"published\")\n",
        "        if agent_dir.endswith(\"/\"+agent.id):\n",
        "          agent.num_paths_sampled = int(os.path.basename(agent_checkpoint).removeprefix(\"checkpoint_\"))\n",
        "          state_dir = os.path.dirname(agent_checkpoint)\n",
        "          paths_df = df_read_from_csv(state_dir, \"paths\")\n",
        "          comps_df = df_read_from_csv(state_dir, \"components\")\n",
        "          def validate_df(df):\n",
        "            assert len(df[\"agent_id\"].unique()) == 1, len(df[\"agent_id\"].unique())\n",
        "            assert df[\"agent_id\"].unique()[0] == agent.id, df[\"agent_id\"].unique()[0]\n",
        "          validate_df(paths_df)\n",
        "          validate_df(comps_df)\n",
        "          Path.paths = []\n",
        "          Path.last_saved = int(paths_df.id.max())\n",
        "          Component.last_saved = int(comps_df.id.max())\n",
        "          Path.counter = Path.last_saved + 1\n",
        "          Component.counter = Component.last_saved + 1\n",
        "          print(\"CONTINUING FROM STATE\", agent.num_paths_sampled)\n",
        "          population_df.append(df_read_from_csv(state_dir, \"population\"))\n",
        "        else:\n",
        "          for _, arow in add_pdf.iterrows():\n",
        "            for cid in arow[\"components\"].split(\",\"):\n",
        "              Population.unmutable_comp2agents[cid].add(arow[\"agent_id\"])\n",
        "        population_df.append(add_pdf)\n",
        "  population_df = pd.concat(population_df)\n",
        "  # Create paths.\n",
        "  cfid2comp = {}\n",
        "  for _, row in population_df.drop_duplicates([\"agent_id\", \"id\"]).iterrows():\n",
        "    hparams = {}\n",
        "    metrics = {}\n",
        "    for k, v in row.iteritems():\n",
        "      if type(v) is float and math.isnan(v):\n",
        "        continue\n",
        "      if k.startswith(\"hparams.\"):\n",
        "        if type(v) == str and (v.startswith(\"{\") or v.startswith(\"[\")):\n",
        "          v = json.loads(v)\n",
        "        hparams[k.removeprefix(\"hparams.\")] = v\n",
        "      if k.startswith(\"metrics.\"):\n",
        "        metrics[k.removeprefix(\"metrics.\")] = v\n",
        "    if \"accounted_params\" in metrics:\n",
        "      del metrics[\"accounted_params\"]  # Triggers recomputation.\n",
        "    comps = []\n",
        "    for c_fid in row[\"components\"].split(\",\"):\n",
        "      if c_fid not in cfid2comp:\n",
        "        c_agent_id, cid = c_fid.split(':')\n",
        "        cfid2comp[c_fid] = Component(name=None, agent_id=c_agent_id, params=None, id=int(cid))\n",
        "      comps.append(cfid2comp[c_fid])\n",
        "    path = Path(hparams, comps, parent=None, agent_id=row[\"agent_id\"],\n",
        "                task_name=row[\"task_name\"], id=int(row[\"id\"]))\n",
        "    path.metrics = metrics\n",
        "    agent.pop.paths[row[\"agent_id\"]].append(path)\n",
        "  # Add components representing subpaths.\n",
        "  for path in [p for paths in agent.pop.paths.values() for p in paths]:\n",
        "    if \"paths\" in path.hparams:\n",
        "      for k in path.hparams[\"paths\"]:\n",
        "        sub_path = agent.pop.get_path_from_full_id(\n",
        "            path.hparams[\"paths\"][k][\"agent_id\"], path.hparams[\"paths\"][k][\"id\"])\n",
        "        path.components.append(ComponentPath(name=k, path=sub_path))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0AyzNaTl002"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m6vSjSIvNPq4"
      },
      "outputs": [],
      "source": [
        "@partial(jax.jit, static_argnames=\"model\")\n",
        "def eval_step(params, inputs, labels, model):\n",
        "  logits = model.apply({\"params\": params}, inputs, train=False)\n",
        "  # Avg accuracy on the batch.\n",
        "  return (logits.argmax(axis=-1) == labels).mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_m2xl8XR7cWy"
      },
      "outputs": [],
      "source": [
        "@partial(jax.jit, static_argnames=[\"model\", \"optimizer\", \"format_params_fn\"], donate_argnums=[0, 2])\n",
        "def train_step(params, fixed_params, opt_state, inputs, labels, model, optimizer, format_params_fn):\n",
        "  def loss_fn(params, fixed_params, inputs, labels):\n",
        "    logits = model.apply(\n",
        "        {\"params\": format_params_fn(merge_params(params, fixed_params))},\n",
        "        inputs, train=True)\n",
        "    labels = jax.nn.one_hot(labels, logits.shape[-1])\n",
        "    return -jnp.mean(jnp.sum(labels * nn.log_softmax(logits), axis=-1))\n",
        "  grads = jax.grad(loss_fn)(params, fixed_params, inputs, labels)\n",
        "  updates, opt_state = optimizer.update(grads, opt_state, params=params)\n",
        "  params = optax.apply_updates(params, updates)\n",
        "  return params, opt_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RrvyCDaw0KFZ"
      },
      "outputs": [],
      "source": [
        "def execute_train_step(path, train_batch):\n",
        "  path.params_device, path.opt_state_device = train_step(\n",
        "      path.params_device,\n",
        "      path.fixed_params_device,\n",
        "      path.opt_state_device,\n",
        "      train_batch[\"image\"],\n",
        "      train_batch[\"label\"],\n",
        "      path.model,\n",
        "      path.optimizer,\n",
        "      path.model_factory.get_comps2model_fn())\n",
        "\n",
        "def execute_eval_step(path, eval_batch):\n",
        "  path.accs.append(\n",
        "      eval_step(\n",
        "          path.model_factory.get_comps2model_fn()(merge_params(\n",
        "              path.params_device, path.fixed_params_device)),\n",
        "          eval_batch[\"image\"],\n",
        "          eval_batch[\"label\"],\n",
        "          path.model))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lRjic_IpGYJU"
      },
      "outputs": [],
      "source": [
        "PREV_LOOP_END = time.time()\n",
        "\n",
        "def train_loop(agent, paths, ds_train, ds_validation, devices):\n",
        "  global PREV_LOOP_END\n",
        "  timing = {}\n",
        "  task = agent.task\n",
        "  for p_id, path in enumerate(paths):\n",
        "    path.device_id = p_id % len(devices)\n",
        "    path.device = devices[path.device_id]\n",
        "    path.optimizer = get_optimizer(\n",
        "        opt_lr=path.hparams[\"opt_lr\"],\n",
        "        opt_lr_schedule=path.hparams[\"opt_lr_schedule\"],\n",
        "        opt_lr_warmup_ratio=path.hparams[\"opt_lr_warmup_ratio\"],\n",
        "        opt_momentum=path.hparams[\"opt_momentum\"],\n",
        "        opt_nesterov=path.hparams[\"opt_nesterov\"],\n",
        "        num_train_batches_between_validations=task.num_train_batches_between_validations,\n",
        "        num_validations_per_path_training=agent.config.num_validations_per_path_training)\n",
        "    path.best_params_local = None\n",
        "    path.best_quality = None\n",
        "    path.best_score = path.parent.score() if path.agent_id == path.parent.agent_id else -np.inf\n",
        "    path.evals = []\n",
        "    path.exe_thread = None\n",
        "  gc.collect()\n",
        "  # Tranfer parameters to devices.\n",
        "  for path in paths:\n",
        "    path.params_device = jax.device_put(path.get_trainable_params(), path.device)\n",
        "    path.fixed_params_device = jax.device_put(path.get_fixed_params(), path.device)\n",
        "    path.opt_state_device = jax.jit(path.optimizer.init, device=path.device)(path.params_device)\n",
        "  iter_ds_validation = iter(ds_validation)\n",
        "  # Train loop.\n",
        "  print((\"\\t\".join([f\"{path.id}\" for path in paths]) + \"\\t\u003c Path id\").expandtabs(8))\n",
        "  for t_step, train_batch in zip(\n",
        "      range(agent.config.num_validations_per_path_training\n",
        "            * task.num_train_batches_between_validations),\n",
        "      ds_train):\n",
        "    if t_step == 0:\n",
        "      timing[\"start_train\"] = time.time()\n",
        "    for path in paths:\n",
        "      train_batch = jax.device_put(train_batch, path.device)\n",
        "      if path.exe_thread is not None:\n",
        "        path.exe_thread.join()\n",
        "      path.exe_thread = Thread(target=execute_train_step, args=(path, train_batch))\n",
        "      path.exe_thread.start()\n",
        "    if t_step == 0:\n",
        "      [p.exe_thread.join() for p in paths]\n",
        "      timing[\"end_train_compile\"] = time.time()\n",
        "    # Evaluation on validation set.\n",
        "    if (t_step+1) % task.num_train_batches_between_validations == 0:\n",
        "      for path in paths:\n",
        "        path.accs = []\n",
        "      for e_step, eval_batch in zip(range(task.num_validation_batches), iter_ds_validation):\n",
        "        if e_step == 0:\n",
        "          start_eval_round = time.time()\n",
        "          if \"start_eval\" not in timing:\n",
        "            timing[\"start_eval\"] = start_eval_round\n",
        "        for path in paths:\n",
        "          eval_batch = jax.device_put(eval_batch, path.device)\n",
        "          path.exe_thread.join()\n",
        "          path.exe_thread = Thread(target=execute_eval_step, args=(path, eval_batch))\n",
        "          path.exe_thread.start()\n",
        "        if e_step == 0 and \"end_eval_compile\" not in timing:\n",
        "          [p.exe_thread.join() for p in paths]\n",
        "          timing[\"end_eval_compile\"] = time.time()\n",
        "      # Get params of best models.\n",
        "      qs = []\n",
        "      eval_idx = (t_step+1) // task.num_train_batches_between_validations\n",
        "      for path in paths:\n",
        "        path.exe_thread.join()\n",
        "        quality = np.mean(path.accs)\n",
        "        del path.accs\n",
        "        qs.append(f\"{quality:.4f}\")\n",
        "        path.evals.append(quality)\n",
        "        # Set quality in metrics for current score computation.\n",
        "        path.metrics[\"quality\"] = quality\n",
        "        path_score = path.score()\n",
        "        if path_score \u003e path.best_score:\n",
        "          path.best_params_local = jax.device_get(path.params_device)\n",
        "          path.best_score = path_score\n",
        "          path.best_quality = quality\n",
        "          qs[-1] += \"*\"\n",
        "      time_train = time.time() - PREV_LOOP_END\n",
        "      avg_path_time = (time_train / eval_idx) / len(paths)\n",
        "      print((\"\\t\".join(qs) + f\"\\t\u003c Eval {eval_idx}\").expandtabs(8),\n",
        "            f\"tot:{time_train:.1f}s\", f\"avg/path:{avg_path_time:.1f}s\")\n",
        "      timing[\"time_eval\"] = timing.get(\"time_eval\", 0) + (time.time() - start_eval_round)\n",
        "      del eval_batch\n",
        "  del train_batch\n",
        "  for path in paths:\n",
        "    del path.params_device\n",
        "    del path.fixed_params_device\n",
        "    del path.opt_state_device\n",
        "    del path.optimizer\n",
        "    del path.exe_thread\n",
        "  gc.collect()\n",
        "  timing[\"end_train\"] = time.time()\n",
        "  time_init = timing[\"start_train\"] - PREV_LOOP_END\n",
        "  time_train_compile = timing[\"end_train_compile\"] - timing[\"start_train\"]\n",
        "  time_eval_compile = timing[\"end_eval_compile\"] - timing[\"start_eval\"]\n",
        "  time_eval = timing[\"time_eval\"] - time_eval_compile\n",
        "  time_train = timing[\"end_train\"] - timing[\"end_train_compile\"] - time_eval - time_eval_compile\n",
        "  PREV_LOOP_END = timing[\"end_train\"]\n",
        "  for path in paths:\n",
        "    path.metrics[\"time_init\"] = time_init\n",
        "    path.metrics[\"time_train_compile\"] = time_train_compile\n",
        "    path.metrics[\"time_eval_compile\"] = time_eval_compile\n",
        "    path.metrics[\"time_train\"] = time_train\n",
        "    path.metrics[\"time_eval\"] = time_eval\n",
        "    path.metrics[\"timestamp_end\"] = PREV_LOOP_END\n",
        "    path.metrics[\"num_params\"] = get_num_params(path.get_all_params())\n",
        "    path.metrics[\"num_trainable_params\"] = get_num_params(path.get_trainable_params())\n",
        "    path.metrics[\"quality\"] = max(path.evals)\n",
        "    path.metrics[\"evals\"] = json.dumps([float(v) for v in path.evals])\n",
        "    if path.best_params_local != None:\n",
        "      path.metrics[\"improved\"] = True\n",
        "      path.update_trainable(path.best_params_local)\n",
        "      assert path.best_quality == path.metrics[\"quality\"]\n",
        "      assert path.best_score == path.score()\n",
        "    else:\n",
        "      path.metrics[\"improved\"] = False\n",
        "      # Sampled path will be dropped if not improved, so skip paramter update.\n",
        "      assert path.best_quality == None\n",
        "    del path.best_params_local\n",
        "    del path.best_score\n",
        "    del path.best_quality\n",
        "    del path.evals\n",
        "  pqs = []\n",
        "  qs = []\n",
        "  psc = []\n",
        "  sc = []\n",
        "  for path in paths:\n",
        "    if path.task_name == path.parent.task_name:\n",
        "      metric_suffix = \"\" if path.agent_id == path.parent.agent_id else \"A\"\n",
        "      pqs.append(f\"{path.parent.metrics['quality']:.4f}{metric_suffix}\")\n",
        "      psc.append(f\"{path.parent.score():.4f}{metric_suffix}\")\n",
        "    else:\n",
        "      pqs.append(\"NEW\")\n",
        "      psc.append(\"NEW\")\n",
        "    qs.append(f\"{path.metrics['quality']:.4f}\")\n",
        "    sc.append(f\"{path.score():.4f}\")\n",
        "    if path.metrics[\"improved\"]:\n",
        "      sc[-1] += \"+\"\n",
        "  print((\"\\t\".join(pqs) + \"\\t\u003c Parent best quality\").expandtabs(8))\n",
        "  print((\"\\t\".join(qs) + \"\\t\u003c Path best quality\").expandtabs(8))\n",
        "  print((\"\\t\".join(psc) + \"\\t\u003c Parent score\").expandtabs(8))\n",
        "  print((\"\\t\".join(sc) + \"\\t\u003c Path score\").expandtabs(8))\n",
        "  print(\"time\\tINIT\\tCOMPtrn\\tCOMPevl\\tTRN\\tEVAL\".expandtabs(8))\n",
        "  print(f\"(s)\\t{time_init:.1f}\\t{time_train_compile:.1f}\\t{time_eval_compile:.1f}\\t{time_train:.1f}\\t{time_eval:.1f}\".expandtabs(8))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JVT8nwIWMAVf"
      },
      "outputs": [],
      "source": [
        "def has_test_quality(path):\n",
        "  return (\"test_quality\" in path.metrics and not math.isnan(path.metrics[\"test_quality\"]))\n",
        "\n",
        "# Run final eval on test set.\n",
        "def run_test_eval(path, task, test_immutability=False):\n",
        "  if not test_immutability and has_test_quality(path):\n",
        "    return  # Skip if test_quality already computed and no immutability test required.\n",
        "  eval_st = time.time()\n",
        "  ds_test = task.get_ds(\"test\", path.hparams)\n",
        "  params = path.get_all_params()\n",
        "  # Running on same device can allow to reuse the fn compiled for validation.\n",
        "  if not hasattr(path, \"device\"):\n",
        "    path.device = random.choice(jax.local_devices())  # Otherwise pick random device.\n",
        "  params_device = jax.device_put(path.model_factory.get_comps2model_fn()(params), path.device)\n",
        "  acc_sum = []\n",
        "  tot_num_samples = 0\n",
        "  # Warning: if repeat() is called on this dataset, then this loop never ends.\n",
        "  for batch in ds_test:\n",
        "    acc_avg = eval_step(params_device, batch[\"image\"], batch[\"label\"], path.model)\n",
        "    batch_size = batch[\"label\"].shape[0]\n",
        "    # Need to recompute sum because last batch can have different size to allow\n",
        "    # for exact eval on the test set.\n",
        "    acc_sum.append(acc_avg * batch_size)\n",
        "    tot_num_samples += batch_size\n",
        "  del params_device\n",
        "  acc_avg = np.sum(acc_sum) / tot_num_samples\n",
        "  if has_test_quality(path):  # Test that new accuracy is close to recorded one.\n",
        "    assert test_immutability\n",
        "    print(f\"Immutability test of path {path.id}: {path.metrics['test_quality']}~={acc_avg}\")\n",
        "    assert np.isclose(path.metrics[\"test_quality\"], acc_avg, rtol=IMMUTABILITY_RELATIVE_TOLLERANCE)\n",
        "  path.metrics[\"test_quality\"] = acc_avg\n",
        "  print(f\"TEST QUALITY: {acc_avg}\\nTEST TIME: {time.time()-eval_st:.2f}s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZubWF7rhmHju"
      },
      "source": [
        "# Main"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VcF5h8pSKx6V"
      },
      "outputs": [],
      "source": [
        "AGENT = get_agent_class(AGENT_CLASS)(system_state_dir=SYSTEM_STATE_DIR, task_name=TASK_NAME, num_paths_sampled_max=NUM_PATHS_SAMPLED_MAX)\n",
        "AGENT.run()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "mu4Net2",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
